pytorch中加载pt文件

在上一篇文章中，我们使用MNIST数据集对建立的卷积神经网络进行了训练，并加载测试集进行测试，最终的识别精度达到了99%。

但是测试流程只是最终给到了整体的测试结果，没有很直观的给到我们看识别效果，所以接下来随机的在MNIST数据集中抽取几张照片，或者自行手写一些数字，输入到训练好的网络中进行识别，来验证结果是不是正确。

如之前的文章所述，使用pytorch 加载的MNIST数据集是以二进制形式存在的，需要进行可视化转换及保存，具体的操作可以参看之前的文章。

张欣In：MNIST数据集详解及可视化处理（pytorch）zhuanlan.zhihu.com

我们在训练集或者测试集中随机抽取五张图片，并以图片内容命名，放到模型和python文件同一文件夹内。

然后使用 torch.load 来加载已经训练完成的模型，依次读取文件夹内的jpg图片，并输出结果观察是否与图片名称一致。

import torch
from mnist import *
import glob
import cv2
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
import torchvision
from skimage import io,transform


if __name__ =='__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.load('./MNIST.pth') #加载模型
    model = model.to(device)
    model.eval（)    #把模型转为test模式
    
    #循环读取文件夹内的jpg图片并输出结果
    for jpgfile in glob.glob(r'./*.jpg'):
        print(jpgfile)#打印图片名称，以与结果进行对照
        img = cv2.imread(jpgfile)  #读取要预测的图片，读入的格式为BGR
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)#图片转为灰度图，因为mnist数据集都是灰度图
        img=np.array(img).astype(np.float32)
        img=np.expand_dims(img,0)
        img=np.expand_dims(img,0)#扩展后，为[1，1，28，28]
        img=torch.from_numpy(img)
        img = img.to(device)
        output=model(Variable(img))
        prob = F.softmax(output, dim=1)
        prob = Variable(prob)
        prob = prob.cpu().numpy()  #用GPU的数据训练的模型保存的参数都是gpu形式的，要显示则先要转回cpu，再转回numpy模式
        print(prob)  #prob是10个分类的概率
        pred = np.argmax(prob) #选出概率最大的一个
        print(pred.item())

运行程序后，会在当前目录中进行遍历，所有jpg格式的文件被依次送入网络，并打印出最终的结果。

.EIGHT.jpg
[[0. 0. 0. 0. 0. 0. 0. 0. 1. 0.]]
8
.FIVE.jpg
[[0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]]
5
.FOUR.jpg
[[0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]]
4
.NINE.jpg
[[0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]]
9
.TWO.jpg
[[0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]]
2

如上，是使用的数据集内的图片进行测试，下面我们手写一张照片，试试结果。

打开windows的画图软件，手写一个数字3，然后进行对数字区域进行裁剪并保存，如下图所示。

在送入网络之前，需要对截取到的数字3进行处理，我们以灰度图的方式读入，然后缩放到28x28大小。

img

因为MNIST数据集中的图片都是黑底白字，所以此处还需要图片进行反色处理。

height,width=img.shape
dst=np.zeros((height,width),np.uint8)
for i in range(height):
    for j in range(width):
        dst[i,j]=255-img[i,j]

img = dst

处理完成后的图片和之前的步骤就一样了，送入网络，输出结果。附上完整代码。

import torch
from mnist import *
import glob
import cv2
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
import torchvision
from skimage import io,transform


if __name__ =='__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.load('./MNIST.pth') #加载模型
    model = model.to(device)
    model.eval（)    #把模型转为test模式

    img = cv2.imread('./THREE.jpg', 0)  #以灰度图的方式读取要预测的图片
    img = cv2.resize(img, (28, 28))

    height,width=img.shape
    dst=np.zeros((height,width),np.uint8)
    for i in range(height):
        for j in range(width):
            dst[i,j]=255-img[i,j]

    img = dst

    img=np.array(img).astype(np.float32)
    img=np.expand_dims(img,0)
    img=np.expand_dims(img,0)#扩展后，为[1，1，28，28]
    img=torch.from_numpy(img)
    img = img.to(device)
    output=model(Variable(img))
    prob = F.softmax(output, dim=1)
    prob = Variable(prob)
    prob = prob.cpu().numpy()  #用GPU的数据训练的模型保存的参数都是gpu形式的，要显示则先要转回cpu，再转回numpy模式
    print(prob)  #prob是10个分类的概率
    pred = np.argmax(prob) #选出概率最大的一个
    print(pred.item())

最终的结果如下：

[[0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]]
3

感兴趣的，可以自行尝试手写并输入其他的数字进行测试。