Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
counterfactual-world-models
/
cwm
/eval
/Segmentation
/archive
/configs
/mask_rcnn_vitdet_b_100ep.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.solver import WarmupParamScheduler | |
from detectron2.modeling.backbone.vit import get_vit_lr_decay_rate | |
import os | |
from ..common.coco_loader_lsj import dataloader | |
from detectron2.data.datasets import register_coco_instances | |
from detectron2.config import CfgNode, LazyConfig | |
# 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_vitdet_v2.py" | |
model = LazyConfig.load(cfg_file).model | |
model.backbone.square_pad = 512 # change input size to 512x512 | |
# 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 = ( | |
#"detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth?matching_heuristics=True" | |
#"/ccn2/u/honglinc/cwm_checkpoints/ablation_3frame_16x16_no_clumping_mr0.98/checkpoint-799-encoder.pth" | |
"/ccn2/u/honglinc/cwm_checkpoints/mae_vitb/mae_pretrain_vit_base-encoder.pth" | |
) | |
train.output_dir = os.path.dirname(train.init_checkpoint) + "/coco_finetune_512_v3" | |
root = os.path.expanduser(os.getenv("DETECTRON2_DATASETS", "datasets")) | |
register_coco_instances("cls_agnostic_coco", {}, | |
os.path.join(root, "coco/annotations/coco_cls_agnostic_instances_val2017.json"), | |
os.path.join(root, "coco/val2017") | |
) | |
dataloader.test.dataset.names = 'cls_agnostic_coco' | |
# 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}} |