mmselfsup/configs/selfsup/_base_/models/mocov3_vit-small-p16.py

39 lines
1.1 KiB
Python

# model settings
temperature = 0.2
model = dict(
type='MoCoV3',
base_momentum=0.99,
backbone=dict(
type='VisionTransformer',
arch='mocov3-small', # embed_dim = 384
img_size=224,
patch_size=16,
stop_grad_conv1=True),
neck=dict(
type='NonLinearNeck',
in_channels=384,
hid_channels=4096,
out_channels=256,
num_layers=3,
with_bias=False,
with_last_bn=True,
with_last_bn_affine=False,
with_last_bias=False,
with_avg_pool=False,
vit_backbone=True),
head=dict(
type='MoCoV3Head',
predictor=dict(
type='NonLinearNeck',
in_channels=256,
hid_channels=4096,
out_channels=256,
num_layers=2,
with_bias=False,
with_last_bn=True,
with_last_bn_affine=False,
with_last_bias=False,
with_avg_pool=False),
temperature=temperature),
loss=dict(type='mmcls.CrossEntropyLoss', loss_weight=2 * temperature))