forked from open-mmlab/mmpretrain
-
Notifications
You must be signed in to change notification settings - Fork 0
/
simclr_resnet50_8xb32-coslr-200e_in1k.py
47 lines (43 loc) · 1.36 KB
/
simclr_resnet50_8xb32-coslr-200e_in1k.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
_base_ = [
'../_base_/datasets/imagenet_bs32_simclr.py',
'../_base_/schedules/imagenet_lars_coslr_200e.py',
'../_base_/default_runtime.py',
]
# model settings
model = dict(
type='SimCLR',
backbone=dict(
type='ResNet',
depth=50,
norm_cfg=dict(type='SyncBN'),
zero_init_residual=True),
neck=dict(
type='NonLinearNeck', # SimCLR non-linear neck
in_channels=2048,
hid_channels=2048,
out_channels=128,
num_layers=2,
with_avg_pool=True),
head=dict(
type='ContrastiveHead',
loss=dict(type='CrossEntropyLoss'),
temperature=0.1),
)
# optimizer
optim_wrapper = dict(
type='OptimWrapper',
optimizer=dict(type='LARS', lr=0.3, momentum=0.9, weight_decay=1e-6),
paramwise_cfg=dict(
custom_keys={
'bn': dict(decay_mult=0, lars_exclude=True),
'bias': dict(decay_mult=0, lars_exclude=True),
# bn layer in ResNet block downsample module
'downsample.1': dict(decay_mult=0, lars_exclude=True),
}))
# runtime settings
default_hooks = dict(
# only keeps the latest 3 checkpoints
checkpoint=dict(type='CheckpointHook', interval=10, max_keep_ckpts=3))
# NOTE: `auto_scale_lr` is for automatically scaling LR
# based on the actual training batch size.
auto_scale_lr = dict(base_batch_size=256)