_base_ = [
    './swin-tiny-upernet_ade-pretrain.py'
]
checkpoint_file = './checkpoints/swin_small_patch4_window7_224_20220317-7ba6d6dd.pth'  # noqa
model = dict(
    backbone=dict(
        init_cfg=dict(type='Pretrained', checkpoint=checkpoint_file),
        depths=[2, 2, 18, 2]),
    decode_head=dict(in_channels=[96, 192, 384, 768], num_classes=150),
    auxiliary_head=dict(in_channels=384, num_classes=150))
training_steps = 80000
param_scheduler = [
    dict(
        type='LinearLR', start_factor=1e-6, by_epoch=False, begin=0, end=8000),
    # dict(
    #     type='PolyLR',
    #     eta_min=0.0,
    #     power=1.0,
    #     begin=10000,
    #     end=training_steps,
    #     by_epoch=False,
    # )
    dict(
        type='CosineAnnealingLR',
        begin=8000,
        end=training_steps,
        eta_min=1e-7,
        by_epoch=False
    )
]

# By default, models are trained on 8 GPUs with 2 images per GPU
train_dataloader = dict(batch_size=16, num_workers=24)
val_dataloader = dict(batch_size=1)
test_dataloader = val_dataloader
train_cfg = dict(max_iters=training_steps, type='IterBasedTrainLoop', val_interval=4000)