def create_model(num_classes=21, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
pre_train_path = "./src/resnet50.pth"
backbone = resnet50()
model = SSD640(backbone=backbone,
num_classes=num_classes,
pretrain_path=pre_train_path)
# # https://ngc.nvidia.com/catalog/models -> search ssd -> download FP32
# pre_ssd_path = "./src/nvidia_ssdpyt_fp32.pt"
# pre_model_dict = torch.load(pre_ssd_path, map_location=device)
# pre_weights_dict = pre_model_dict["model"]
#
# del_conf_loc_dict = {}
# for k, v in pre_weights_dict.items():
# split_key = k.split(".")
# if "conf" in split_key:
# continue
# del_conf_loc_dict.update({k: v})
#
# missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict, strict=False)
# if len(missing_keys) != 0 or len(unexpected_keys) != 0:
# print("missing_keys: ", missing_keys)
# print("unexpected_keys: ", unexpected_keys)
return model
def __init__(self, pretrain_path=None):
super(Backbone, self).__init__()
net = resnet50()
self.out_channels = [1024, 512, 512, 256, 256, 256]
if pretrain_path is not None:
net.load_state_dict(torch.load(pretrain_path))
self.feature_extractor = nn.Sequential(*list(net.children())[:7])
conv4_block1 = self.feature_extractor[-1][0]
# 修改conv4_block1的步距，从2->1
conv4_block1.conv1.stride = (1, 1)
conv4_block1.conv2.stride = (1, 1)
conv4_block1.downsample[0].stride = (1, 1)
