def alexnet(pretrained=False, **kwargs): # 224*224
if pretrained:
model = AlexNet(**kwargs)
pretrained_state_dict = torch.load(
'./Authority/alexnet-owt-4df8aa71.pth')
now_state_dict = model.state_dict() # 返回model模块的字典
pretrained_state_dict.pop('classifier.6.weight')
pretrained_state_dict.pop('classifier.6.bias')
now_state_dict.update(pretrained_state_dict)
model.load_state_dict(now_state_dict)
return model
return AlexNet(**kwargs)
from torch.autograd import Variable
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision as tv
import torchvision.transforms as transforms
# Tensor的读取与保存
a = torch.Tensor(3, 4)
a.cuda()
torch.save(a, 'a.pth')
b = torch.load('a.pth')
c = torch.load('a.pth', map_location=lambda sto, loc: sto)
# ----------------------------------------------------------
torch.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import AlexNet
model = AlexNet()
model.state_dict().keys()
# model的保存与加载
torch.save(model.state_dict(), 'alexnet.pth')
model.load_state_dict(torch.load('alexnet.pth'))
opt = torch.optim.Adam(model.parameters(), lr=0.1)
# 优化器的参数读取与保存
torch.save(opt.state_dict(), 'opt.pth')
opt.load_state_dict(torch.load('opt.pth'))
