使用ResNet对CIFAR10数据集进行识别分类（pytorch）_cifar10数据集resnet分类

LeNet5的结构比较简单，分类准确率只有50%左右。
ResNet属于中等规模复杂度的网络，性能比LeNet5会强大不少。
本例采用最简单的ResNet18模型，实现对CIFAR数据集的10分类。

还是按照之前的流程，分四步完成网络的搭建和训练。
编程过程中发现最好还是使用模块化编程，不然容易写bug出来。

代码如下：

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
import torch.optim as optim
from torch.nn import functional as F
import matplotlib.pyplot as plt
import time

# # 使用ResNet18网络训练CIFAR10数据集实现10分类

start = time.time()
# Step 1 : prepare dataset

batch_size = 32

cifar_train = datasets.CIFAR10("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)  # 导入训练数据集 添加三个变换 第一个将图片裁剪至32*32大小；第二个将格式转变成tensor；第三个使用均值归一化，使数据均匀分布在0-1之间

cifar_train_loader = DataLoader(cifar_train, batch_size=batch_size, shuffle=True,)  # 做打乱处理

cifar_test = datasets.CIFAR10("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_loader = DataLoader(cifar_test, batch_size=batch_size, shuffle=False, )  # 与上面相同，但测试集不需要打乱


# Step2: design model

# 先定义Res模块 res模块是残差神经网络中的残差运算单元
class ResBlock(nn.Module):  # 同样继承至nn.Module
    def __init__(self, ch_in, ch_out, stride=1):
        super(ResBlock, self).__init__()

        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)  # stride对图片尺寸的大小有重要的影响
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)  # 两个卷积层 两个batchnorm
        # shortcut 短接层
        self.extra = nn.Sequential()
        if ch_out != ch_in:
            # let [b, ch_in, h, w] ----> [b, ch_out, h, w]
            self.extra = nn.Sequential(
                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),  # 此处的stride设置与conv1一样 保证大小一致 可以相加
                nn.BatchNorm2d(ch_out)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        # extra shortcut
        out = self.extra(x) + out
        out = F.relu(out)
        return out

# 再定义ResNet类

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=3, padding=0),
            nn.BatchNorm2d(64)
        )
        # follow 4 block
        self.resblock = nn.Sequential(
            ResBlock(64, 128, stride=2),
            ResBlock(128, 256, stride=2),
            ResBlock(256, 512, stride=2),
            ResBlock(512, 512, stride=2),  # 512是经验之谈 一般channel提升到512  同时图片尺寸需要减少
        )

        self.outlayer = nn.Linear(512, 10)
        # 总的结构为：1个卷积[b, 3, 32, 32]-->[b, 64, 32, 32]+4个残差块[b, 64, 32, 32]-->[b, 512, 2, 2]+1个全连接层[b, 512]-->[b, 10]
        # 4个残差块后还有一个全局池化的操作，实现[b, 512, 2, 2]-->[b, 512, 1, 1],并[b, 512, 1, 1]-->[b, 512*1*1]

    def forward(self, x):
        x = F.relu(self.conv1(x))
        # [b, 64, h, w] ----> [b, 512, h, w]
        x = self.resblock(x)

        # print("after conv: ", x.shape)  # [b, 512, 2, 2]
        x = F.adaptive_max_pool2d(x, [1, 1])  # [b, 512, h, w] ---> [b, 512, 1, 1] 不管w,h是多少，都变成1*1的(均值)
        # print("after pool: ", x.shape)
        # flatten to 1 dim
        x = x.view(x.size(0), -1)
        x = self.outlayer(x)
        return x


model = ResNet()
print(model)  # 打印模型结构
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)  # 放到GPU上


# Step3: construct Loss and Optimizer

criterion = torch.nn.CrossEntropyLoss()  # 分类一般使用交叉熵
optimizer = optim.Adam(model.parameters(), lr=0.001)


# Step4: Train and Test

def train(epoch):
    running_loss = 0
    model.train()  # 设置为train模式
    for batch_idx, (x, label) in enumerate(cifar_train_loader, 0):
        x, label = x.to(device), label.to(device)
        optimizer.zero_grad()

        # forward
        outputs = model(x)
        loss = criterion(outputs, label)
        # backward
        loss.backward()
        # update
        optimizer.step()

    print("Epoch: ", epoch, "Loss is: ", loss.item())

def test(epoch):
    correct = 0
    total = 0
    model.eval()  # 设置为test模式
    with torch.no_grad():  # 以下内容不需要构建计算图，不需要计算梯度 这一句可加可不加
        for data in cifar_test_loader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            pred = outputs.argmax(dim=1)
            total += labels.size(0)  # 每次循环都把这一批的batch_size加上，就得到总的数量
            correct += torch.eq(pred, labels).float().sum().item()  # 对比预测和label相同的数量 即为预测正确的数量
    print("Epoch", epoch, "Accuracy on test set: %d %%" % (100 * correct / total))
    return correct / total



if __name__ == "__main__":
    epoch_list = []
    acc_list = []
    for epoch in range(50):
        train(epoch)
        acc = test(epoch)
        epoch_list.append(epoch)
        acc_list.append(acc)

    plt.plot(epoch_list, acc_list)
    plt.xlabel("Epoch")
    plt.ylabel("Acc")
    plt.grid()
    plt.show()
    
    end = time.time()
    print("Total Time: ", end - start)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

结果如图：

可以看出准确率大约在72%左右，相比LeNet5上升了20多个百分点。
训练50轮，总时间为2303秒。