Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 2,327 Bytes
6dfcb0f |
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 48 49 50 51 52 53 54 55 56 57 58 59 |
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
from detectron2.data.datasets import register_coco_instances
# 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 = (
# "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth?matching_heuristics=True"
'/ccn2/u/honglinc/cwm_checkpoints/ablation_3frame_16x16_no_clumping_mr0.90/checkpoint-799-encoder.pth'
)
train.output_dir = os.path.dirname(train.init_checkpoint) + "/coco_finetune_512_v3"
train.eval_period = 1e9
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}} |