Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
counterfactual-world-models
/
cwm
/eval
/Segmentation
/archive
/configs
/mask_rcnn_vitdet_b_100ep_dino.py
from functools import partial | |
from fvcore.common.param_scheduler import MultiStepParamScheduler | |
from detectron2 import model_zoo | |
from detectron2.config import LazyCall as L | |
from detectron2.config import CfgNode, LazyConfig | |
from detectron2.solver import WarmupParamScheduler | |
from detectron2.modeling.backbone.vit import get_vit_lr_decay_rate | |
import os | |
from ..common.coco_loader_lsj import dataloader | |
# model = model_zoo.get_config("common/models/mask_rcnn_vitdet.py").model | |
# model.backbone.square_pad = 512 # change input size to 512x512 | |
cfg_file = "./models/mask_rcnn_cwm.py" | |
model = LazyConfig.load(cfg_file).model | |
# Initialization and trainer settings | |
train = model_zoo.get_config("common/train.py").train | |
train.amp.enabled = True | |
train.ddp.fp16_compression = True | |
train.init_checkpoint = ( | |
'/home/honglinc/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth' | |
) | |
train.output_dir = '/ccn2/u/honglinc/cwm_checkpoints/dinov2_coco_finetune_512' | |
# model.backbone.net.window_size = 0 | |
# model.backbone.net.window_block_indexes = [] | |
# model.backbone.net.use_rel_pos = False | |
# model.backbone.net.drop_path_rate = 0. | |
# Schedule | |
# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep | |
# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep | |
# train.max_iter = 184375 | |
# milestones = [163889, 177546] | |
# 50 ep = 30730 iters * 96 images/iter / 118000 images/ep | |
train.max_iter = 61458 | |
milestones = [54629, 59182] | |
lr_multiplier = L(WarmupParamScheduler)( | |
scheduler=L(MultiStepParamScheduler)( | |
values=[1.0, 0.1, 0.01], | |
milestones=milestones, | |
num_updates=train.max_iter, | |
), | |
warmup_length=250 / train.max_iter, | |
warmup_factor=0.001, | |
) | |
# Optimizer | |
optimizer = model_zoo.get_config("common/optim.py").AdamW | |
optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.7) | |
optimizer.params.overrides = {"pos_embed": {"weight_decay": 0.0}} |