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'))