def init_model(cfg):
model_cfg = edict()
model_cfg.crop_size = (320, 480)
model_cfg.input_normalization = {
'mean': [.485, .456, .406],
'std': [.229, .224, .225]
# 'mean': [.838, .855, .770],
# 'std': [.281, .210, .300]
}
model_cfg.num_max_points = 12
model_cfg.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(model_cfg.input_normalization['mean'],
model_cfg.input_normalization['std']),
])
model = get_deeplab_model(backbone='resnet50',
deeplab_ch=128,
aspp_dropout=0.20,
norm_radius=180)
model.to(cfg.device)
model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0))
model.feature_extractor.load_pretrained_weights()
return model, model_cfg
def load_deeplab_is_model(state_dict,
device,
backbone='auto',
deeplab_ch=128,
aspp_dropout=0.2,
cpu_dist_maps=False,
norm_radius=260):
if backbone == 'auto':
num_backbone_params = len([
x for x in state_dict.keys() if 'feature_extractor.backbone' in x
and not ('num_batches_tracked' in x)
])
if num_backbone_params <= 181:
backbone = 'resnet34'
elif num_backbone_params <= 276:
backbone = 'resnet50'
elif num_backbone_params <= 531:
backbone = 'resnet101'
else:
raise NotImplementedError('Unknown backbone')
if 'aspp_dropout' in state_dict:
aspp_dropout = float(state_dict['aspp_dropout'].cpu().numpy())
else:
aspp_project_weight = [
v for k, v in state_dict.items()
if 'aspp.project.0.weight' in k
][0]
deeplab_ch = aspp_project_weight.size(0)
if deeplab_ch == 256:
aspp_dropout = 0.5
model = get_deeplab_model(backbone=backbone,
deeplab_ch=deeplab_ch,
aspp_dropout=aspp_dropout,
cpu_dist_maps=cpu_dist_maps,
norm_radius=norm_radius)
model.load_state_dict(state_dict, strict=False)
for param in model.parameters():
param.requires_grad = False
model.to(device)
model.eval()
return model