Convolutional neural network code (CIFAR image classification) and basic framework based on Pytorch

1. Notice

1. The data set used in this code is CIFAR10, which can be downloaded and loaded through the following code segment.
Import torchvision needs to be referenced

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. The network does not support the situation that the size of each picture in the data set is different from each other. If you build a data set yourself or load another data set, please make a unified format for the size of the data set first. It is recommended to change to 3 * 32 * 32. If you change to other formats, calculate the total number of pixels after nn.Flatten() and replace 1024 in nn.Linear(1024, 10) 3. The test set results do not output the final

category Judgment, only supports the output of the correct rate (correct number/total number of test sets)

4. Support tensorboard

5. Add every 100 iterations to calculate the time difference

6. The activation function is not added, you need to add it yourself

7. Due to the relatively simple basic framework, the performance of the model is slightly worse. When running 165epoch, the test set achieved the highest accuracy rate of 68.5%.

2. Network Model Framework

The basic architecture idea is to

read data→build minibacth→select GPU or CPU training→select loss function→build forward transfer network→select GSD model to descend and set hyperparameters→start iteration→calculate loss function→backpropagation→update parameters → output result → test

3. Complete code

It can be run directly on the code of kaggle, and the data set can be selected as 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()