mmpretrain/configs/simsiam/simsiam_resnet50_8xb32-cosl...

53 lines
1.4 KiB
Python

_base_ = [
'../_base_/datasets/imagenet_bs32_mocov2.py',
'../_base_/schedules/imagenet_sgd_coslr_200e.py',
'../_base_/default_runtime.py',
]
# model settings
model = dict(
type='SimSiam',
backbone=dict(
type='ResNet',
depth=50,
norm_cfg=dict(type='SyncBN'),
zero_init_residual=True),
neck=dict(
type='NonLinearNeck',
in_channels=2048,
hid_channels=2048,
out_channels=2048,
num_layers=3,
with_last_bn_affine=False,
with_avg_pool=True),
head=dict(
type='LatentPredictHead',
loss=dict(type='CosineSimilarityLoss'),
predictor=dict(
type='NonLinearNeck',
in_channels=2048,
hid_channels=512,
out_channels=2048,
with_avg_pool=False,
with_last_bn=False,
with_last_bias=True)),
)
# optimizer
# set base learning rate
lr = 0.05
optim_wrapper = dict(
type='OptimWrapper',
optimizer=dict(type='SGD', lr=lr, weight_decay=1e-4, momentum=0.9),
paramwise_cfg=dict(custom_keys={'predictor': dict(fix_lr=True)}))
# runtime settings
default_hooks = dict(
# only keeps the latest 3 checkpoints
checkpoint=dict(type='CheckpointHook', interval=10, max_keep_ckpts=3))
# additional hooks
custom_hooks = [
dict(type='SimSiamHook', priority='HIGH', fix_pred_lr=True, lr=lr)
]