pytorch官方demo实例（搭建Lenet并基于CIFAR10训练）

pytorch官方demo实例（搭建Lenet并基于CIFAR10训练）

demo教程来自官网 Training a Classifier

model.py文件

import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):#搭建模型第一步,定义类LeNet,需继承nn.module副类
    def __init__(self):#初始化函数，实现搭建模型需要的网络层结构
        super(LeNet, self).__init__()#使用super函数，解决在多重继承中调用副类中出现的问题
        self.conv1 = nn.Conv2d(3, 16, 5)#nn.Conv2d定义卷积层 第一参数代表输入特征层的参数，16个卷积核，卷积核的尺寸5x5
        self.pool1 = nn.MaxPool2d(2, 2)#池化核大小为2 x 2 步距也为2 经过最大池化将高度和宽度缩减为原来一半，池化层只改变特征矩阵的高和宽
        self.conv2 = nn.Conv2d(16, 32, 5)#定义第二个卷积层，输入特征层深度为16，32个卷积核，卷积深度为5x5
        self.pool2 = nn.MaxPool2d(2, 2)#同上
        self.fc1 = nn.Linear(32*5*5, 120)#全连接层，输入为一维向量，需要将得到的特征矩阵展平，第一层节点个数为120
        self.fc2 = nn.Linear(120, 84)#第二层84个节点
        self.fc3 = nn.Linear(84, 10)#10需要根据训练的数据集进行修改
#经卷积后的矩阵尺寸大小计算公式为 N=(W-F+2P)/S+1
#输入图片大小W x W Filter大小F x F 步长S padding的像素数P
    def forward(self, x):#定义正向传播过程，x代表输入的数据batch,channel,height,width
        x = F.relu(self.conv1(x))    # input(3, 32, 32) output(16, 28, 28)#W=32 卷积核大小5 F=5 padding=0 S=1 所以输出28
        x = self.pool1(x)            # output(16, 14, 14)
        x = F.relu(self.conv2(x))    # output(32, 10, 10)#（14-5+0）/1+1=10 输出为32
        x = self.pool2(x)            # output(32, 5, 5)
        x = x.view(-1, 32*5*5)       # output(32*5*5)#和全连接层拼接，第一维度为batch，再展平数据
        x = F.relu(self.fc1(x))      # output(120)
        x = F.relu(self.fc2(x))      # output(84)
        x = self.fc3(x)              # output(10)
        return x

import  torch
input1 = torch.rand([32,3,32,32])
model = LeNet()
print(model)
output = model(input1)

train.py

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

from model import LeNet


def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),#Resize将图像缩放为32 x 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')#通过PIL载入图像
    im = transform(im)  # [C, H, W]#将图像转化为tensor，将图像标准化
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]#增加一个维度，pytorch要求四个参数 batch channel height weight 而输入的图片没有batch需要增加第一个维度

    with torch.no_grad():#不需要求损失梯度
        outputs = net(im)#将图像传入到网络中
        predict = torch.max(outputs, dim=1)[1].data.numpy()
    print(classes[int(predict)])#将index传入到classes寻找预测的类


if __name__ == '__main__':
    main()

predict.py

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

from model import LeNet


def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),#Resize将图像缩放为32 x 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')#通过PIL载入图像
    im = transform(im)  # [C, H, W]#将图像转化为tensor，将图像标准化
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]#增加一个维度，pytorch要求四个参数 batch channel height weight 而输入的图片没有batch需要增加第一个维度

    with torch.no_grad():#不需要求损失梯度
        outputs = net(im)#将图像传入到网络中
        predict = torch.max(outputs, dim=1)[1].data.numpy()
    print(classes[int(predict)])#将index传入到classes寻找预测的类


if __name__ == '__main__':
    main()