_base_ = [
    '../_base_/models/upernet_beit.py',
    '../_base_/datasets/my_dataset_model.py',
    '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_40k_check_4000.py',
]

norm_cfg = dict(
    type='BN',
)

crop_size = (512, 512)

data_preprocessor = dict(
    size=(512, 512),
    mean=[
        94.94709810464303,
        61.72942233949928,
        75.93763705236906,
    ],
    std=[
        44.005506081132594,
        42.69595666984776,
        44.99354156225523,
    ],
    bgr_to_rgb=False,
)

model = dict(
    backbone=dict(
        type='BEiT',
        img_size=(512, 512),
        embed_dims=1024,
        num_layers=24,
        num_heads=16,
        init_values=1e-06,
        drop_path_rate=0.2,
        out_indices=[
            7,
            11,
            15,
            23,
        ],
    ),
    data_preprocessor=dict(
        size=(512, 512),
        mean=[
            94.94709810464303,
            61.72942233949928,
            75.93763705236906,
        ],
        std=[
            44.005506081132594,
            42.69595666984776,
            44.99354156225523,
        ],
        bgr_to_rgb=False,
    ),
    pretrained='./My_Local_Model/pretrain/beit_large_patch16_224_pt22k_ft22k.pth',
    neck=dict(
        type='Feature2Pyramid',
        embed_dim=1024,
        rescales=[
            4,
            2,
            1,
            0.5,
        ],
    ),
    decode_head=dict(
        in_channels=[
            1024,
            1024,
            1024,
            1024,
        ],
        num_classes=36,
        channels=1024,
        norm_cfg=dict(
            type='BN',
        ),
        loss_decode=dict(
            type='DiceLoss',
            use_sigmoid=False,
            loss_weight=1.0,
        ),
    ),
    auxiliary_head=dict(
        norm_cfg=dict(
            type='BN',
        ),
        in_channels=1024,
        num_classes=36,
        loss_decode=dict(
            type='DiceLoss',
            use_sigmoid=False,
            loss_weight=0.4,
        ),
    ),
)

optim_wrapper = dict(
    _delete_=True,
    type='AmpOptimWrapper',
    optimizer=dict(
        type='AdamW',
        lr=2e-05,
        betas=(0.9, 0.999),
        weight_decay=0.05,
    ),
    constructor='LayerDecayOptimizerConstructor',
    paramwise_cfg=dict(
        num_layers=24,
        layer_decay_rate=0.95,
    ),
    accumulative_counts=2,
)

param_scheduler = [
    dict(
        type='LinearLR',
        start_factor=1e-06,
        by_epoch=False,
        begin=0,
        end=1500,
    ),
    dict(
        type='PolyLR',
        power=0.9,
        begin=1500,
        end=40000,
        eta_min=1e-05,
        by_epoch=False,
    ),
]

train_dataloader = dict(
    batch_size=2,
)

