Kaggle猫狗大战——基于Pytorch的CNN网络分类：预测模型结果（4）
Kaggle猫狗大战——基于Pytorch的CNN网络分类：预测模型结果（4）
本文是Kaggle猫狗大战项目的最后一步了，写一个predict.py，在命令行输入随便找的猫狗图片，使用训练好的模型进行预测。这块比较简单，就话不多说，直接上代码。
predict.py
import sys
import torch
import os
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from PIL import Image
def predict(dir):
model = torch.load('best_model.pt') # 载入模型
model.train(False) # 模型设定为train(false)模式
test_transforms = transforms.Compose([
transforms.Resize(224), # 图像短边长度变为input_size
transforms.CenterCrop(224), # 从正中间剪正方形
transforms.ToTensor(),
transforms.Normalize([0.4864, 0.4533, 0.4154], [0.2625, 0.2558, 0.2586])
])
img_PIL = Image.open(dir) # 打开文件路径
image_tensor = test_transforms(img_PIL) # 将图片文件转换为tensor文件
image_tensor.unsqueeze_(0)
image_tensor = image_tensor.to(device) # 由CPU.tensor文件转换为GPU.tensor文件
out = model(image_tensor) # 得到预测结果，并且从大到小排序
_,indices = torch.sort(out, descending=True) # 返回每个预测的百分数
print(['The given picture is '+(class_names[idx]) for idx in indices[0][:1]])
if __name__ == '__main__':
order = sys.argv[1] # 从命令行获取图片名
test_dir = ('../test/'+order) # 组合图片路径
class_names = ['cat','dog']
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
predict(test_dir)
unsqueeze函数
预测程序的代码非常简单，也很好理解，我就都已标注的形式写在了文件里。唯一让我比较迷惑的就是image_tensor.unsqueeze_(0)这一步，没有会报错，所以特意去查了一下。
unsqueeze()函数和squeeze()函数是一对函数，作用是对张量Tensor的维度进行压缩或者扩充。我对tensor的理解就是pytorch中的tensor张量与numpy中的数组相似，因此其维度也很好理解。
例如：a=torch.IntTensor([[1,2,3],[4,5,6]])其维度就是[2,3]。而b = torch.IntTensor([[[1,2,3],[4,5,6]]])其维度是[1,2,3]。（先是一个11的数组，然后是一个23的数组。）正是因为张量b其第0个维度是一维，所以可以对张量b进行压缩，将其压缩成[2,3]的张量。（注意，压缩只能对一维的维度进行，但是膨胀可以随便加在哪儿）。
c=torch.squeeze(b,0)
这样一来，张量c的维度就是[2,3]。同样地，也可以对a进行膨胀操作：
d=torch.unsqueeze(a,0)
如此一来，张量d的维度就是[1,2,3]。
注意，也可以对后面的维度进行操作，比如，e=torch.unsqueeze(a,1)，则张量e的维度就是[2,1,3]。
e=([[[0,1,2]],[[3,4,5]]])
下面则是我个人对这里应用该函数的一些看法，因为我们在data.py里面进行数据打包时，每个数据包中都包含了多张图片的数据，而这里是对单张图片进行处理，所以需要在单张图片的张量外再套一层，使其与data.py中的数据包类型相同。
命令行操作
我们的目标是：
（1）将你想识别的图片存到和data同文件夹的test文件夹中；
（2）打开cmd；
（3）在cmd中输入python predict.py XXX.jpg；（XXX.jpg是你要识别的图片的名字）
（4）cmd中回复：The given picture is cat/dog
操作非常简单，利用order = sys.argv[1]指令就能实现从cmd中读取代码，但是在实操时需要注意一些地方，好多博客中用到了sys.argv，但是直接打开cmd往往无法运行，这是由于位置与环境不对。
第一步，到程序所在文件夹里，在搜索栏（此电脑那一栏）删去其他的，输入cmd，回车；
这样一来，跳出来的命令行界面就是这样的：
第二步，如果你使用了虚拟环境（基本都会用），激活虚拟环境
第三步，输入python predict.py XXX.jpg，前面的python predict.py是调用程序的，后面的XXX.jpg使用sys.argv可以直接读取成字符串。
