用CNN对CIFAR10进行分类（pytorch）

时间：2023-03-09 06:50:06
用CNN对CIFAR10进行分类（pytorch）

CIFAR10有60000个\(32*32\)大小的有颜色的图像，一共10种类别，每种类别有6000个。

训练集一共50000个图像，测试集一共10000个图像。

先载入数据集

import numpy as np

import torch

import torch.optim as optim

from torchvision import datasets

import torchvision.transforms as transforms

transform = transforms.Compose([

    transforms.ToTensor(),

    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

])

trainset = datasets.CIFAR10(root='./data', train=True,

                                        download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,

                                          shuffle=True, num_workers=2)

testset = datasets.CIFAR10(root='./data', train=False,

                                       download=True, transform=transform)

testloader = torch.utils.data.DataLoader(testset, batch_size=4,

                                         shuffle=False, num_workers=2)

再定义网络架构

import torch.nn.functional as F

import torch.nn as nn

class classifier(nn.Module):

    def __init__(self):

        super().__init__()

        '''输入为3*32*32，尺寸减半是因为池化层'''

        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)   #输出为16*16*16

        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)  #输出为32*8*8

        self.pool = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(32 * 8 * 8, 512)

        self.fc2 = nn.Linear(512, 10)

        self.dropout = nn.Dropout(0.2)     #防止过拟合

    def forward(self, x):

        x = self.pool(F.relu(self.conv1(x)))

        x = self.pool(F.relu(self.conv2(x)))

        x = x.view(-1, 32 * 8 * 8)

        x = self.dropout(x)

        x = F.relu(self.fc1(x))

        x = self.dropout(x)

        x = self.fc2(x)

        return x

开始训练！

model = classifier()

criterion = nn.CrossEntropyLoss()

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

epochs = 10

for e in range(epochs):

    train_loss = 0

    for data, target in train_loader:

        optimizer.zero_grad()

        output = model(data)

        loss = criterion(output, target)

        loss.backward()

        optimizer.step()

        train_loss += loss.item() * data.size(0)	#loss.item()是平均损失，平均损失*batch_size=一次训练的损失

    train_loss = train_loss/len(train_loader.dataset)

    print('Epoch: {} \t Training Loss:{:.6f}'.format(e+1, train_loss))
下面是损失的输出
Epoch: 1 	 Training Loss:1.366521

Epoch: 2 	 Training Loss:1.063830

Epoch: 3 	 Training Loss:0.916826

Epoch: 4 	 Training Loss:0.799573

Epoch: 5 	 Training Loss:0.708303

Epoch: 6 	 Training Loss:0.627443

Epoch: 7 	 Training Loss:0.564043

Epoch: 8 	 Training Loss:0.503542

Epoch: 9 	 Training Loss:0.465513

Epoch: 10 	 Training Loss:0.418729

看看在验证集上的表现如何！

class_correct = list(0. for i in range(10))

class_total = list(0. for i in range(10))

with torch.no_grad():

    for data in testloader:

        images, labels = data

        outputs = model(images)

        _, predicted = torch.max(outputs, 1)

        c = (predicted == labels).squeeze()

        for i in range(4):

            label = labels[i]

            class_correct[label] += c[i].item()

            class_total[label] += 1

for i in range(10):

    print('Accuracy of %5s : %2d %%' % (

        classes[i], 100 * class_correct[i] / class_total[i]))
以及它的输出
Accuracy of plane : 74 %

Accuracy of   car : 76 %

Accuracy of  bird : 55 %

Accuracy of   cat : 56 %

Accuracy of  deer : 54 %

Accuracy of   dog : 54 %

Accuracy of  frog : 81 %

Accuracy of horse : 72 %

Accuracy of  ship : 74 %

Accuracy of truck : 68 %

相关文章