学习笔记1 pytorch实现Lenet5

参考资料：https://cuijiahua.com/blog/2018/01/dl_3.html

代码实现：　

import torch
from torch import nn
from torch.nn import functional as F
class Lenet5(nn.Module):
    """
    for cifar10 dataset
    """
    def __init__(self):
        super(Lenet5, self).__init__()
        self.conv_unit=nn.Sequential(
            #卷积层 x:[b,3,32,32] =>[b,6,30,30]
            nn.Conv2d(3,6,kernel_size=3,stride=1,padding=0),
            #池化层 [b,6,30,30]=>[b,6,15,15]
            nn.MaxPool2d(kernel_size=2,stride=2,padding=0),
            #卷积层 [b,6,15,15] =>[b,16,13,13]
            nn.Conv2d(6,16,kernel_size=3,stride=1,padding=0),
            #池化层 [b,16,13,13]=>[b,16,6,6]
            nn.MaxPool2d(kernel_size=2,stride=2,padding=0)
        )
        #flatten
        #全连接层
        self.fc_unit=nn.Sequential(
            nn.Linear(16*6*6,120),
            nn.ReLU(inplace=True),
            nn.Linear(120,84),
            nn.ReLU(inplace=True),
            nn.Linear(84,10)
        )

        #use Cross Entropy loss
        #self.criteon=nn.CrossEntropyLoss()

    def forward(self,x):
        """
        x:[b,3,32,32]
        :param x:
        :return:
        """
        batchsz=x.size(0)
        #[b,3,32,32]=>[b,16,6,6]
        x=self.conv_unit(x)
        #[b,16,6,6] => [b,16*6*6]
        x=x.view(batchsz,16*6*6)
        #[b,16*6*6]=>[b,10]
        logits=self.fc_unit(x)

        return logits
        #
    
#测试函数
def main():
    net=Lenet5()
    tmp = torch.randn(2, 3, 32, 32)
    #通过测试，确定conv_unit输入维度
    out = net(tmp)
    print('conv out:', out.shape)
if __name__=='__main__':
    main()

import  torch
from    torch.utils.data import DataLoader
from    torchvision import datasets
from    torchvision import transforms
from    torch import nn, optim

from    lenet5 import Lenet5

def main():
    #batchsize大小
    batchsz = 32
    #读取一张图片并进行数据增强
    cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor()
    ]), download=True)
    #读取bathsize大小的一批图像
    cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)

    cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor()
    ]), download=True)
    cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)

    #测试图像读取是否正确
    #x, label = iter(cifar_train).next()
    #print('x:', x.shape, 'label:', label.shape)

    #选用gpu设备
    device = torch.device('cuda:0,1')
    model = Lenet5().to(device)
    #model = ResNet18().to(device)
    
    #定义交叉熵评估损失值
    criteon = nn.CrossEntropyLoss().to(device)
    #定义优化器
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    #打印模型信息
    print(model)

    #train
    for epoch in range(1000):
        #model.train()
        for batchidx, (x, label) in enumerate(cifar_train):
            # [b, 3, 32, 32]
            # [b]
            x, label = x.to(device), label.to(device)
            
            #预测值
            logits = model(x)
            # logits: [b, 10]
            # label:  [b]
            # loss: tensor scalar
            loss = criteon(logits, label)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(epoch, 'loss:', loss.item())
        #test
        model.eval()
        with torch.no_grad():
            # test
            total_correct = 0
            total_num = 0
            for x, label in cifar_test:
                # [b, 3, 32, 32]
                # [b]
                x, label = x.to(device), label.to(device)

                # [b, 10]
                logits = model(x)
                # [b] 选取所有类别得分的最大值作为预测类
                pred = logits.argmax(dim=1)
                # [b] vs [b] => scalar tensor
                #计算每个batch正确的预测数量
                correct = torch.eq(pred, label).float().sum().item()
                total_correct += correct
                total_num += x.size(0)
            acc = total_correct / total_num
            print(epoch, 'acc:', acc)

if __name__ == '__main__':
    main()

结果：