network.py 文件源码

def load_pretrained_model(faster_rcnn_model, model_name='vgg16'):
    if model_name == 'vgg16':
        model = models.vgg16(pretrained=True)
        faster_rcnn_model.rpn.features = model.features
        mod = list(model.classifier.children())[:-1]
        faster_rcnn_model.fcs = nn.Sequential(*mod)
    elif model_name == 'resnet101':
        model = models.resnet101(pretrained=True)
        faster_rcnn_model.rpn.features = nn.Sequential(model.conv1, model.bn1, model.relu, model.maxpool,
                                                       model.layer1, model.layer2, model.layer3, model.layer4,
                                                       model.avgpool)
        faster_rcnn_model.fcs = model.fc