python pytorch_Python:PyTorch 保存和加载训练过的网络 (八十)

保存和加载模型

在这个 notebook 中，我将为你展示如何使用 Pytorch 来保存和加载模型。这个步骤十分重要，因为你一定希望能够加载预先训练好的模型来进行预测，或是根据新数据继续训练。

%matplotlib inline

%config InlineBackend.figure_format = 'retina'

import matplotlib.pyplot as plt

import torch

from torch import nn

from torch import optim

import torch.nn.functional as F

from torch.autograd import Variable

from torchvision import datasets, transforms

import helper

# Define a transform to normalize the data

transform = transforms.Compose([transforms.ToTensor(),

transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# Download and load the training data

trainset = datasets.FashionMNIST('F_MNIST_data/', download=True, train=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

# Download and load the test data

testset = datasets.FashionMNIST('F_MNIST_data/', download=True, train=False, transform=transform)

testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=True)

在这里我们可以看见一张图片。

image, label = next(iter(trainloader))

helper.imshow(image[0,:]);

构建网络

在这里，我将使用与第五部分中一样的模型。

class Network(nn.Module):

def __init__(self, input_size, output_size, hidden_layers, drop_p=0.5):

''' Builds a feedforward network with arbitrary hidden layers.

Arguments

---------

input_size: integer, size of the input layer

output_size: integer, size of the output layer

hidden_layers: list of integers, the sizes of the hidden layers

'''

super().__init__()

# Input to a hidden layer

self.hidden_layers = nn.ModuleList([nn.Linear(input_size, hidden_layers[0])])

# Add a variable number of more hidden layers

layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])

self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in layer_sizes])

self.output = nn.Linear(hidden_layers[-1], output_size)

self.dropout = nn.Dropout(p=drop_p)

def forward(self, x):

''' Forward pass through the network, returns the output logits '''

for each in self.hidden_layers:

x = F.relu(each(x))

x = self.dropout(x)

x = self.output(x)

return F.log_softmax(x, dim=1)

训练网络

并使用之前一样的方法来训练网络。

# Create the network, define the criterion and optimizer

model = Network(784, 10, [500, 100])

criterion = nn.NLLLoss()

optimizer &