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()