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)
def get_model(device=None):
# 加载CNN模型
model = AlexNet(num_classes=2)
model.load_state_dict(
torch.load('./models/best_linear_svm_alexnet_car.pth'))
model.eval()
# 取消梯度追踪
for param in model.parameters():
param.requires_grad = False
if device:
model = model.to(device)
return model
def alexnet(num_classes, num_domains=None, pretrained=True):
"""AlexNet model architecture from the
`"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = AlexNet()
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
print('Load pre trained model')
num_ftrs = model.classifier[-1].in_features
model.classifier[-1] = nn.Linear(num_ftrs, num_classes)
nn.init.xavier_uniform_(model.classifier[-1].weight, .1)
nn.init.constant_(model.classifier[-1].bias, 0.)
return model
class Feature_AlexNet(nn.Module):
def __init__(self):
super(Feature_AlexNet, self).__init__()
self.model = AlexNet()
self.model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
self.fc1 = nn.Linear(256 * 6 * 6, 4096)
self.bn1_fc = nn.BatchNorm1d(4096)
self.fc2 = nn.Linear(4096, 2048)
self.bn2_fc = nn.BatchNorm1d(2048)
self.relu = nn.ReLU(inplace=True)
def forward(self, x, reverse=False):
x = self.model.avgpool(self.model.features(x))
x_feat = x.view(x.size(0), 256 * 6 * 6)
x = self.relu(self.bn1_fc(self.fc1(x_feat)))
x = F.dropout(x, training=self.training)
if reverse:
x = grad_reverse(x, self.lambd)
x = self.relu(self.bn2_fc(self.fc2(x)))
return x, x_feat
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'))
