pytorch实现lenet5,训练集CIFAR10

Lenet网络模型

 

 Lenet核心代码

'''
案例 1
'''
import torch
import torch.nn as nn


#leNet 5 
class Lenet5(nn.Module):
    def __init__(self):
        super(Lenet5,self).__init__()

        self.conv_unit = nn.Sequential(
            # x [b,3,32,32] => [b,6,28,28]
            nn.Conv2d(3,6,kernel_size=5,stride=1,padding=0),
            nn.AvgPool2d(kernel_size=2,stride=2,padding=0), #长和宽各变为原来的1半
            #
            nn.Conv2d(6,16,kernel_size=5,stride=1,padding=0),
            nn.AvgPool2d(kernel_size=2,stride=2,padding=0)
        )
        # flatten 在forward里实现
        #fc unit
        self.fc_unit = nn.Sequential(
            nn.Linear(16*5*5,120),
            nn.ReLU(),
            nn.Linear(120,84),
            nn.ReLU(),
            nn.Linear(84,10)
        )
        #测试
        '''#[B,16,5,5] 打平
        tmp = torch.randn(2,3,32,32)
        out = self.conv_unit(tmp)
        print(out.shape)
        '''
        #交叉熵  分类问题用交叉熵，逼近regresssion用mse
        self.criteon = nn.CrossEntropyLoss()
        #pytorch自动根据你写的forward 自己弄Backward
    def forward(self,x):
            batchsz = x.size(0) #x.size返回一个List [b,3,32,32]
            # [b,3,32,32] - > [b,16,5,5]
            x = self.conv_unit(x)
            # [b,16,5,5] - >  [b,16*5*5]
            x = x.view(batchsz,16*5*5) #设置好了第一维，其它内容自动归到第二个维度
            # [b,16*5*5] - > [b,10]
            logits = self.fc_unit(x)
            #[b,10]
            #pred = F.softmax(logits,dim=1)
            #.CrossEntropy 包含softmax 和 error
            #F是函数不需要初始化，nn需要初始化 且nn内部定义了__call__ 可以自动调用forward
            return logits

if __name__ == '__main__':
#测试代码    
    net = Lenet5()
    tmp = torch.randn(2,3,32,32)
    out = net(tmp)
    print('lenet out:{}'.format(out.shape))

import torch.optim as optim
from torchvision import datasets,transforms
from torch.utils.data import random_split,DataLoader
import torch.nn as nn
import torch.nn.functional as F
import torch
from pytorch_39_Lenet import Lenet5 #导入lenet
batch_size1 = 32

def main():
    
    #一次加载一张
    cifar_train = datasets.CIFAR10('cifar',True,transform=transforms.Compose([
        transforms.Resize((32,32)),
        transforms.ToTensor()
    ]),download=True)
    #多线程，批量读取
    cifar_train = DataLoader(cifar_train,batch_size=batch_size1,shuffle =True)
    
    cifar_test = datasets.CIFAR10('cifar',True,transform=transforms.Compose([
        transforms.Resize((32,32)),
        transforms.ToTensor()
    ]),download=True)
    #多线程，批量读取
    cifar_test = DataLoader(cifar_test,batch_size=batch_size1,shuffle =True)

    x,label = iter(cifar_train).next()
    print('x:',x.shape,'label:',label.shape)

    device = torch.device('cuda')
    model = Lenet5().to(device)
    criteon = nn.CrossEntropyLoss().to(device)
    optimizer = optim.Adam(model.parameters(),lr=1e-3)
    print(model)
    model.train()
    for epoch in range(1000):
        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 = criteon(logits,label)
            #backprop
            optimizer.zero_grad() #不清零就是累加
            loss.backward()
            optimizer.step()
        
        print(epoch,loss.item())

    # test
    ''''
    无论requires_grad是Ture还是False，反向传播时都不会自动求导。不会自动构建图
    test可以实现一定速度的提升，并节省一半的显存，因为其不需要保
    存梯度。
    因为会有一些dropout等操作，所以在训练时候需要设置model.train()
    测试的时候只要设置model.eval()
    '''
    model.eval()
    with torch.no_grad():
        total_correct = 0
        total_num = 0
        for x, label in cifar_test:
            #[b,3,32,32]
            #[b]
            x,label = x.to(0), label.to(device)

            #[b,10]
            logits = model(x)
            #[b] vs [b] -> scalar tensor
            pred = logits.argmax(dim=1)
            total_correct += torch.eq(pred,label).float().sum()
            total_num += x.size(0)
        acc = total_correct / total_num
        print(epoch,acc)
main()