🍨 本文为🔗365天深度学习训练营中的学习记录博客
🍦 参考文章：365天深度学习训练营-第P2周：彩色识别
🍖 原作者：K同学啊 | 接辅导、项目定制
🚀 文章来源：K同学的学习圈子

环境

系统：Linux
语言： Python 3.8.10
深度学习框架：PyTorch 2.0.0+cu118

步骤

环境设置

包引用

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transformsimport numpy as np
import matplotlib.pyplot as plt
from torchinfo import summary # 方便像tensorflow一样打印模型

硬件设备

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

数据准备

数据集下载与加载

train_dataset = datasets.CIFAR10(root='data', train=True, download=True, transform=transforms.ToTensor()) # 不要忘记这个transform
test_dataset = datasets.CIFAR10(root='data', train=False, download=True, transform=transforms.ToTensor())

数据集预览

image, label = train_dataset[0]
print(image.shape)
plt.figure(figsize=(20,4))
for i in range(20):image, label = train_dataset[i]plt.subplot(2, 10, i+1)plt.imshow(image.numpy().transpose(1,2,0)plt.axis('off')plt.title(label) # 加载的数据集没有对应的名称，暂时展示它们的id

数据集预览

数据集准备

batch_size = 32
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

模型设计

class Model(nn.Module):def __init__(self, num_classes):super().__init__()# 3x3的卷积无padding每次宽高-2# 2x2的最大池化，每次宽高缩短为原来的一半# 32x32 -> conv1 -> 30x30 -> maxpool -> 15x15self.conv1 = nn.Conv2d(3, 64, kernel_size=3)# 15x15 -> conv2 -> 13x13 -> maxpool -> 6x6self.conv2 = nn.Conv2d(64, 64, kernel_size=3)# 6x6 -> conv3 -> 4x4 -> maxpool -> 2x2self.conv3 = nn.Conv2d(64, 128, kernel_size=3)self.maxpool = nn.MaxPool2d(2),self.flatten = nn.Flatten(),self.fc1 = nn.Linear(2*2*128, 256)self.fc2 = nn.Linear(256, num_classes)def forward(self, x):x = F.relu(self.conv1(x))x = self.maxpool(x)x = F.relu(self.conv2(x))x = self.maxpool(x)x = F.relu(self.conv3(x))x = self.maxpool(x)x = self.flatten(x)x = F.relu(self.fc1(x))x = self.fc2(x)return xmodel = Model(10).to(device)
summary(model, input_size=(1, 3, 32, 32))

模型结构图

模型训练

超参数设置

learning_rate = 1e-2
epochs = 10
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

helper函数

def train(train_loader, model, loss_fn, optimizer):size = len(train_loader.dataset)num_batches = len(train_loader)train_loss, train_acc = 0, 0for x, y in train_loader:x, y = x.to(device), y.to(device)preds = model(x)loss = loss_fn(preds, y)optimizer.zero_grad()loss.backward()optimizer.step()train_loss += loss.item()train_acc += (preds.argmax(1) == y).type(torch.float).sum().item()train_loss /= num_batchestrain_acc /= sizereturn train_loss, train_accdef test(test_loader, model, loss_fn):size = len(test_loader.dataset)num_batches = len(test_loader)test_loss, test_acc = 0, 0with torch.no_grad():for x, y in test_loader:x, y = x.to(device), y.to(device)preds = model(x)loss = loss_fn(preds, y)test_loss += loss.item()test_acc += (preds.argmax(1) == y).type(torch.float).sum().item()test_loss /= num_batchestest_acc /= sizereturn test_loss, test_accdef fit(train_loader, test_loader, model, loss_fn, optimizer, epochs):train_loss, train_acc = [], []test_loss, test_acc = [], []for epoch in range(epochs):model.train()epoch_train_loss, epoch_train_acc = train(train_loader, model, loss_fn, optimizer)model.eval()epoch_test_loss, epoch_test_acc = test(test_loader, model, loss_fn)train_loss.append(epoch_train_loss)train_acc.append(epoch_train_acc)test_loss.append(epoch_test_loss)test_acc.append(epoch_test_acc)return train_loss, train_acc, test_loss, test_acc

正式训练

train_loss, train_acc, test_loss, test_acc = fit(train_loader, test_loader, model, loss_fn, optimizer, 20)

训练结果

结果呈现

series = range(len(train_loss))
plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(series, train_loss, label='train loss')
plt.plot(series, test_loss, label='validation loss')
plt.legend(loc='upper right')
plt.title('Loss')
plt.subplot(1,2,2)
plt.plot(series, train_acc, label='train accuracy')
plt.plot(series, test_acc, label='validation accuracy')
plt.legend(loc='lower right')
plt.title('Accuracy')