base_glean 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
                            LoggerHook, ParamSchedulerHook)
from mmengine.model import MMSeparateDistributedDataParallel
from mmengine.optim import CosineAnnealingLR, OptimWrapper
from mmengine.runner import IterBasedTrainLoop

from mmagic.engine.runner import MultiTestLoop, MultiValLoop
from mmagic.evaluation import MAE, PSNR, SSIM

[文档]_base_ = '../default_runtime.py'

# DistributedDataParallel
[文档]model_wrapper_cfg = dict(
    type=MMSeparateDistributedDataParallel, find_unused_parameters=True)

[文档]save_dir = './work_dirs'

[文档]val_evaluator = [
    dict(type=MAE),
    dict(type=PSNR),
    dict(type=SSIM),
]
[文档]test_evaluator = val_evaluator

[文档]train_cfg = dict(type=IterBasedTrainLoop, max_iters=300_000, val_interval=5000)
[文档]val_cfg = dict(type=MultiValLoop)
[文档]test_cfg = dict(type=MultiTestLoop)

# optimizer
[文档]optim_wrapper = dict(
    constructor='MultiOptimWrapperConstructor',
    generator=dict(
        type=OptimWrapper,
        optimizer=dict(type='Adam', lr=1e-4, betas=(0.9, 0.99))),
    discriminator=dict(
        type=OptimWrapper,
        optimizer=dict(type='Adam', lr=1e-4, betas=(0.9, 0.99))),
)

# learning policy
[文档]param_scheduler = dict(
    type=CosineAnnealingLR, by_epoch=False, T_max=600_000, eta_min=1e-7)

[文档]default_hooks = dict(
    checkpoint=dict(
        type=CheckpointHook,
        interval=5000,
        save_optimizer=True,
        by_epoch=False,
        out_dir=save_dir,
        save_best=['MAE', 'PSNR', 'SSIM'],
        rule=['less', 'greater', 'greater']),
    timer=dict(type=IterTimerHook),
    logger=dict(type=LoggerHook, interval=100),
    param_scheduler=dict(type=ParamSchedulerHook),
    sampler_seed=dict(type=DistSamplerSeedHook),
)