pytorch学习笔记——模型的保存与载入（12）

训练模型的结果得保存起来才能得到迭代的改进吖！我们使用了torch.save和torch.load方法对原先的模型进行保存与载入。代码如下：

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
from sklearn.datasets import load_digits
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import torch as t
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt
from graphviz import Digraph
import torch
from torch.autograd import Variable
from torchsummary import summary
from PIL import Image
from torchvision.transforms import transforms
import random

class train_mini_mnist(t.utils.data.Dataset):
    def __init__(self):
        self.X,self.y=load_digits(return_X_y=True)
        self.X=self.X
        self.X_train,self.X_test,self.y_train,self.y_test=train_test_split(self.X,self.y,random_state=0)

    def __getitem__(self, index):
            img, target = np.array(self.X_train[index].reshape(8,8),dtype=int), int(self.y_train[index])
            plt.imshow(img)
            img=transforms.ToPILImage()(img)
            img=img.rotate(random.randint(-20,20))#填充白色
            img=transforms.ToTensor()(img)
            return img/15.,target

    def __len__(self):
        return len(self.y_train)

class test_mini_mnist(t.utils.data.Dataset):
    def __init__(self):
        self.X,self.y=load_digits(return_X_y=True)
        self.X=self.X/15.
        self.X_train,self.X_test,self.y_train,self.y_test=train_test_split(self.X,self.y,random_state=0)
    def __getitem__(self, index):
        return t.tensor(self.X_test[index].reshape(1,8,8),dtype=torch.float32),self.y_test[index]
    def __len__(self):
        return len(self.y_test)

BATCH_SIZE=8
LEARNING_RATE=3e-3
EPOCHES=100

train_data=train_mini_mnist()
test_data=test_mini_mnist()

train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE, shuffle=True)

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Sequential(#(1, 8, 8)
            nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=1),#(8, 8, 8) 
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),#(16, 8, 8) 
            nn.BatchNorm2d(16),
            nn.ReLU(),#(16,8,8)
            nn.MaxPool2d(kernel_size=2)#(16,4,4)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),#(32, 4, 4) 
            nn.BatchNorm2d(32),
            nn.ReLU(),#(32,4,4)
            nn.MaxPool2d(kernel_size=2)#(32,2,2)
        )
        self.fc = nn.Sequential(
            nn.Linear(32*2*2, 64),
            nn.Dropout(0.5),
            nn.Linear(64, 10)
            )
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)#相当于Flatten
        x = self.fc(x)
        return x

def eval_on_dataloader(name,loader,len):
    acc = 0.0
    with torch.no_grad():
        for data in loader:
            images, labels = data
            outputs = net(images)
            predict_y = torch.max(outputs, dim=1)[1]#torch.max返回两个数值，一个是最大值，一个是最大值的下标
            acc += (predict_y == labels).sum().item()
        accurate = acc / len
        return accurate

def plot_train_and_test_result(train_accs,test_accs):
    epoches=np.arange(1,len(train_accs)+1,dtype=np.int32)
    plt.plot(epoches,train_accs,label="train_accuracy")
    plt.plot(epoches,test_accs,label="test_accuracy")
    plt.xlabel('epoches')
    plt.ylabel('accuracy')
    plt.legend()



#net = SimpleNet()
net=torch.load('net20.pkl')
print(net.parameters)
summary(net, (1, 8, 8))

loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(net.parameters(), lr = LEARNING_RATE,weight_decay=3e-4)

for name,parameters in net.named_parameters():
    print(name,":",parameters.size())
best_acc = 0.0

train_accs,test_accs=[],[]
for epoch in range(EPOCHES):

    net.train()#切换到训练模式

    for step, data in enumerate(train_loader, start=0):
        images, labels = data
        optim.zero_grad()#将优化器的梯度清零
        logits = net(images)#网络推断的输出
        loss = loss_fn(logits, labels.long())#计算损失函数
        loss.backward()#反向传播求梯度
        optim.step()#优化器进一步优化

        rate = (step+1)/len(train_loader)
        a = "*" * int(rate * 50)
        b = "." * int((1 - rate) * 50)
        print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate*100), a, b, loss), end="")
    print()

    net.eval()#切换到测试模式

    train_acc=eval_on_dataloader("train",train_loader,train_data.__len__())
    test_acc=eval_on_dataloader("test",test_loader,test_data.__len__())
    train_accs.append(train_acc)
    test_accs.append(test_acc)
    print("epoch:",epoch,"train_acc:",train_acc," test_acc:",test_acc)
    if test_acc>best_acc:
        best_acc=test_acc
        torch.save(net, 'net'+str(best_acc)+'.pkl')

print('Finished Training')
plot_train_and_test_result(train_accs,test_accs)
torch.save(net, 'net20.pkl')
plt.show()

经过多次微调，我们的模型达到了99.56%的最终准确率.（我知道调整的策略不是很科学，大佬轻喷~）