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