def __init__(self, name: str, dilation: bool, freeze_bn: bool, last_stage_block=14):
super().__init__()
if "resnet" in name:
norm_layer = FrozenBatchNorm2d if freeze_bn else nn.BatchNorm2d
# here is different from the original DETR because we use feature from block3
self.body = getattr(resnet_module, name)(
replace_stride_with_dilation=[False, dilation, False],
pretrained=is_main_process(), norm_layer=norm_layer, last_layer='layer3')
self.num_channels = 256 if name in ('resnet18', 'resnet34') else 1024
elif "RepVGG" in name:
print("#" * 10 + " Warning: Dilation is not valid in current code " + "#" * 10)
repvgg_func = get_RepVGG_func_by_name(name)
self.body = repvgg_func(deploy=False, last_layer="stage3", freeze_bn=freeze_bn,
last_stage_block=last_stage_block)
self.num_channels = 192 # 256x0.75=192
elif "LightTrack" in name:
print("#" * 10 + " Warning: Dilation is not valid in current code " + "#" * 10)
path_backbone = [[0], [4, 5], [0, 2, 5, 1], [4, 0, 4, 4], [5, 2, 1, 0]]
ops = (3, 2)
self.body = build_subnet(path_backbone, ops)
if is_main_process():
print(self.body)
self.num_channels = 96
else:
raise ValueError("Unsupported net type")
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool, freeze_bn: bool):
norm_layer = FrozenBatchNorm2d if freeze_bn else nn.BatchNorm2d
# here is different from the original DETR because we use feature from block3
backbone = getattr(resnet_module, name)(
replace_stride_with_dilation=[False, dilation, False],
pretrained=is_main_process(),
norm_layer=norm_layer,
last_layer='layer3')
num_channels = 256 if name in ('resnet18', 'resnet34') else 1024
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers)
def __init__(self, img_sz, model_name, train_backbone):
super().__init__()
if model_name == "vit_deit_base_distilled_patch16_384":
ckpt_name = "https://dl.fbaipublicfiles.com/deit/deit_base_distilled_patch16_384-d0272ac0.pth"
model = deit(img_sz, pretrained=is_main_process(), model_name=model_name, ckpt_name=ckpt_name)
self.num_channels = model.num_features
elif model_name == "swin_base_patch4_window12_384_S16":
model = build_swint(model_name)
self.num_channels = model.num_features[-1]
else:
raise ValueError("Unsupported model name")
if not train_backbone:
for name, parameter in model.named_parameters():
parameter.requires_grad_(False)
self.body = model
def build_backbone_x_cnn(cfg, phase='train'):
"""Without positional embedding, standard tensor input"""
train_backbone = cfg.TRAIN.BACKBONE_MULTIPLIER > 0
backbone = Backbone(cfg.MODEL.BACKBONE.TYPE, cfg.MODEL.BACKBONE.DILATION, cfg.TRAIN.FREEZE_BACKBONE_BN,
cfg.MODEL.BACKBONE.LAST_STAGE_BLOCK)
if phase is 'train':
"""load pretrained backbone weights"""
ckpt_path = None
if hasattr(cfg, "ckpt_dir"):
if cfg.MODEL.BACKBONE.TYPE == "RepVGG-A0":
filename = "RepVGG-A0-train.pth"
elif cfg.MODEL.BACKBONE.TYPE == "LightTrack":
filename = "LightTrackM.pth"
else:
raise ValueError("The checkpoint file for backbone type %s is not found" % cfg.MODEL.BACKBONE.TYPE)
ckpt_path = os.path.join(cfg.ckpt_dir, filename)
if ckpt_path is not None:
print("Loading pretrained backbone weights from %s" % ckpt_path)
ckpt = torch.load(ckpt_path, map_location='cpu')
if cfg.MODEL.BACKBONE.TYPE == "LightTrack":
ckpt, ckpt_new = ckpt["state_dict"], {}
for k, v in ckpt.items():
if k.startswith("features."):
k_new = k.replace("features.", "")
ckpt_new[k_new] = v
ckpt = ckpt_new
missing_keys, unexpected_keys = backbone.body.load_state_dict(ckpt, strict=False)
if is_main_process():
print("missing keys:", missing_keys)
print("unexpected keys:", unexpected_keys)
"""freeze some layers"""
if cfg.MODEL.BACKBONE.TYPE != "LightTrack":
trained_layers = cfg.TRAIN.BACKBONE_TRAINED_LAYERS
# defrost parameters of layers in trained_layers
for name, parameter in backbone.body.named_parameters():
parameter.requires_grad_(False)
if train_backbone:
for trained_name in trained_layers: # such as 'layer2' in layer2.conv1.weight
if trained_name in name:
parameter.requires_grad_(True)
break
return backbone
def __init__(self,
name: str,
train_backbone: bool,
return_interm_layers: bool,
dilation: bool,
freeze_bn: bool,
ckpt_path=None):
if "resnet" in name:
norm_layer = FrozenBatchNorm2d if freeze_bn else nn.BatchNorm2d
# here is different from the original DETR because we use feature from block3
backbone = getattr(resnet_module, name)(
replace_stride_with_dilation=[False, dilation, False],
pretrained=is_main_process(),
norm_layer=norm_layer,
last_layer='layer3')
num_channels = 256 if name in ('resnet18', 'resnet34') else 1024
net_type = "resnet"
elif "RepVGG" in name:
print(
"#" * 10 +
" Freeze_BN and Dilation are not supported in current code "
+ "#" * 10)
# here is different from the original DETR because we use feature from block3
repvgg_func = get_RepVGG_func_by_name(name)
backbone = repvgg_func(deploy=False, last_layer="stage3")
num_channels = 192 # 256x0.75=192
try:
ckpt = torch.load(ckpt_path, map_location='cpu')
missing_keys, unexpected_keys = backbone.load_state_dict(
ckpt, strict=False)
print("missing keys:", missing_keys)
print("unexpected keys:", unexpected_keys)
except:
print("Warning: Pretrained RepVGG weights are not loaded")
net_type = "repvgg"
else:
raise ValueError()
super().__init__(backbone,
train_backbone,
num_channels,
return_interm_layers,
net_type=net_type)
def get_optimizer_scheduler(net, cfg):
train_cls = getattr(cfg.TRAIN, "TRAIN_CLS", False)
if train_cls:
print("Only training classification head. Learnable parameters are shown below.")
param_dicts = [
{"params": [p for n, p in net.named_parameters() if "cls" in n and p.requires_grad]}
]
for n, p in net.named_parameters():
if "cls" not in n:
p.requires_grad = False
else:
print(n)
else:
param_dicts = [
{"params": [p for n, p in net.named_parameters() if "backbone" not in n and p.requires_grad]},
{
"params": [p for n, p in net.named_parameters() if "backbone" in n and p.requires_grad],
"lr": cfg.TRAIN.LR * cfg.TRAIN.BACKBONE_MULTIPLIER,
},
]
if is_main_process():
print("Learnable parameters are shown below.")
for n, p in net.named_parameters():
if p.requires_grad:
print(n)
if cfg.TRAIN.OPTIMIZER == "ADAMW":
optimizer = torch.optim.AdamW(param_dicts, lr=cfg.TRAIN.LR,
weight_decay=cfg.TRAIN.WEIGHT_DECAY)
else:
raise ValueError("Unsupported Optimizer")
if cfg.TRAIN.SCHEDULER.TYPE == 'step':
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.TRAIN.LR_DROP_EPOCH)
elif cfg.TRAIN.SCHEDULER.TYPE == "Mstep":
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=cfg.TRAIN.SCHEDULER.MILESTONES,
gamma=cfg.TRAIN.SCHEDULER.GAMMA)
else:
raise ValueError("Unsupported scheduler")
return optimizer, lr_scheduler