Pytorch 实现Lenet5，并在cifar10数据集上进行训练

拼搏现实的明天。 2022-01-21 09:41 229阅读 0赞

lenet5.py

import  torch
    from    torch import nn
    
    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, 16, 5, 5]
                nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
                nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
                nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
                nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            )
    
            self.fc_unit = nn.Sequential(
                nn.Linear(16*5*5, 120),
                nn.ReLU(),
                nn.Linear(120, 84),
                nn.ReLU(),
                nn.Linear(84, 10)
            )
    
    
            # [b, 3, 32, 32]
            test = torch.randn(2, 3, 32, 32)
            out = self.conv_unit(test)
            # [b, 16, 5, 5]
            print('conv out:', out.shape)
    
    
        def forward(self, x):
            """
            :param x: [b, 3, 32, 32]
            :return:
            """
            batchsz = x.size(0)
            # [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]
            return logits
    
    def main():
    
        net = Lenet5()
    
        test = torch.randn(2, 3, 32, 32)
        out = net(test)
        print("lenet out: ", out)
    
    if __name__ == '__main__':
        main()

main.py

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():
    
        batchsz = 128
    
        cifar_train = datasets.CIFAR10(root='cifar', train=True, transform=transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ]), download=True)
    
        cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
    
        cifar_test = datasets.CIFAR10(root='cifar', train=False, transform=transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ]), 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)
    
        device = torch.device("cuda")
        model = Lenet5().to(device)
        criteon = nn.CrossEntropyLoss().to(device)
        optimizer = optim.Adam(model.parameters(), lr=1e-3)
        print(model)
        for epoch in range(1000):
            model.train()
            for batchidx, (x, label) in enumerate(cifar_train):
                # x: [b, 3, 32, 32]
                # label: [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()
    
            # loss.item()转换为numpy
            print("epoch {0}: loss {1}".format(epoch, loss.item()))
    
            model.eval()
            with torch.no_grad():
                # test
                total_correct = 0
                total_num = 0
                for x, label in cifar_test:
                    # x: [b, 3, 32, 32]
                    # label: [b]
                    x, label = x.to(device), label.to(device)
    
                    # logits: [b, 10]
                    logits = model(x)
                    # [b]
                    pred = logits.argmax(dim=1)
                    total_correct = total_correct + torch.eq(pred, label).float().sum().item()
                    total_num = total_num + x.size(0)
    
                acc = total_correct / total_num
                print("epoch {0}: acc {1}".format(epoch, acc))
    
    
    if __name__ == '__main__':
        main()