加载指定层的预训练模型代码

import torch
import os
import torchvision.models as models
import torch.nn as nn

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device".format(device))


class Net(nn.Module):
    def __init__(self, model):
        super(Net, self).__init__()
        self.resnet50 = nn.Sequential(*list(model.children())[:-1])
        self.fc = nn.Linear(2048, 2)

    def forward(self, x):
        x = self.resnet50(x)
        x = self.fc(x)
        return x


# 先加载在创建net
model = models.resnet50(pretrained=True)
model_weight_path = "./resnet50.pth"
# 等价于if not，条件为假时执行
assert os.path.exists(model_weight_path), "file {} does not exist.".format(model_weight_path)

model.load_state_dict(torch.load(model_weight_path, map_location=device))
net = Net(model)
print(net)

参考：
https://blog.csdn.net/whut_ldz/article/details/78845947
https://blog.csdn.net/xys430381_1/article/details/107038723