使用 PyTorch 搭建网络 - predict_py篇

predict_py篇

python中采用驼峰书写法且首字母大写的变量符号一般表示类名。

学习网络步骤：看原论文+看别人对原论文的理解，学习网络结构，看损失函数计算，看数据集，看别人写的代码，复现代码。

经历以上步骤我们便可以选择合适的框架复现代码，这里使用PyTorch复现网络结构。我们用PyTorch搭建网络可以分为以下几个module，数据处理dataloader.py，网络模型model.py，训练模块train.py，工具模块utils.py，预测模块predict.py。接下来我们将以LeNet为例进行讲解。

由于LeNet使用了CIFAR10数据集，所以我们直接用内置方法生成dataset便可。

目录如下：

• predict_py 综述
• torch.nn.Module.load_state_dic() 方法
• torch.softmax() 方法
• 源码

---

predict_py 综述

在该模块中进行预测，我们需要创建net，并载入train_py中训练好的权重文件，配置好net的parameters。将test中图片经过dataloader_py后转成我们需要的tensor，将该tensor传入net中，进行FP传播得到最终的输出。

torch.nn.Module.load_state_dic() 方法

net = LeNet()
net.load_state_dict(torch.load("Lenet.pth"))

我们创建网络后，使用net对象的load_state_dic()和torch.load()方法载入预训练权重。

torch.softmax() 方法

with torch.no_grad():
    outputs = net(im)
    predict = torch.softmax(outputs, dim = 1)
print(predict)

在全连接层的最后一层接入torch.softmax()方法进行多分裂预测。

源码

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

transforms = transforms.Compose(
    [transforms.Resize((32, 32)),
    transforms.ToTensor(), 
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)

classes = ("plane", "car", "bird", "cat",
            "deer", "dog", "frog", "horse", "ship", "truck")

net = LeNet()
net.load_state_dict(torch.load("Lenet.pth"))

im = Image.open("1.jpg")
im = transforms(im) # [C, H, W]
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]

with torch.no_grad():
    outputs = net(im)
    predict = torch.softmax(outputs, dim = 1)
print(predict)