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
Пример #2
0
    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)