基于Pytorch的卷积神经网络代码(CIFAR图像分类)及基本构架

一、须知

1.本代码所用数据集为CIFAR10，可通过以下代码段进行下载并加载
需要引用 import torchvision

train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())

2.网络不支持数据集中各图片尺寸相互不一的情况，若自行构建数据集或加载别的数据集，请先对数据集尺寸做成统一格式。推荐更改为3 * 32 * 32 ,若更改为其他格式，自行计算nn.Flatten()之后的像素总数，并替换掉nn.Linear(1024, 10)中的1024

3.测试集结果不输出最终类别判断，仅支持正确率（正确个数/测试集总个数） 的输出

4.支持tensorboard

5.加入每100次迭代计算时间差

6.未加入激活函数，需要自行添加

7.由于基础框架比较简单，模型表现效果略差，运行165epoch时，测试集取得最高准确率68.5%

二、网络模型框架

基本构架思路为

读取数据→构建minibacth→选择GPU或CPU训练→选择损失函数→构建前向传递网络→选择GSD模型进行下降并设置超参→开始迭代→计算损失函数→反向传播→更新参数→输出结果→测试

三、完整代码

可直接在kaggle的code上运行，数据集选择cifar10-python即可

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time


train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())

train_dataloader = DataLoader(train_data, batch_size=64, drop_last=True)
test_dataloader = DataLoader(test_data, batch_size=64, drop_last=True)
# print(len(train_dataloader)) #781
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


test_data_size = len(test_dataloader) * 64
print(f'测试集大小为：{test_data_size}')
writer = SummaryWriter("../model_logs")

loss_fn = nn.CrossEntropyLoss(reduction='mean')
loss_fn = loss_fn.to(device)
time_able = False # True

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),# 182528
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)

        return x

model = Model()
model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
epoch = 50
running_loss = 0
total_train_step = 0
total_test_step = 0
if time_able:
    str_time = time.time()
for i in range(epoch):
    print(f'第{i + 1}次epoch')
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = model(imgs)
        loss = loss_fn(output, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 100 == 0:
            if time_able:
                end_time = time.time()
                print(f'{str_time-end_time}')
            print(f'第{total_train_step}次训练，loss = {loss.item()}')
            writer.add_scalar("train_loss", loss.item(), total_train_step)
    # 测试
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = model(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss = total_test_loss + loss
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
    total_test_loss = total_test_loss / test_data_size
    print(f'整体测试集上的loss = {total_test_loss}')
    print(f'整体测试集正确率 = {total_accuracy / test_data_size}')
    writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
    total_test_step += 1

writer.close()