CIFAR10彩色图片识别

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍖 原作者：K同学啊

CIFAR10彩色图片识别

这是我参加训练营的第二周

数据处理

好多项目的数据处理部分思路是相同的。

带入库函数

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision

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

device

加上下面的这一句，，matplotlib 绘制的图形会直接嵌入在 Notebook 的输出单元格中显示，而不是弹出独立窗口。

%matplotlib inline

在我的电脑上pytorch和matplotlib容易冲突，加上下面的三行才能在pytorch环境中运行matplotlib

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["OMP_NUM_THREADS"] = "4"          # 限制线程数
os.environ["MKL_NUM_THREADS"] = "4"

下载数据集

train_ds = torchvision.datasets.CIFAR10('data', 
                                      train=True, 
                                      transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
                                      download=True)

test_ds  = torchvision.datasets.CIFAR10('data', 
                                      train=False, 
                                      transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
                                      download=True)

下载好数据集，得取数据集，先用torch.utils.data.DataLoader取出，以32张图片为一组，取完所有的图片为一轮。训练集要取很多轮，所以要设置shuffle=True，每次取完一轮，顺序不一样。测试集不用取很多轮，所以不用设置。

batch_size=32
train_dl=torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True)
test_dl=torch.utils.data.DataLoader(test_ds, batch_size=batch_size)

训练集有图片和标签，分别赋值给imgs，labels

imgs,labels=next(iter(train_dl))
imgs.shape

看看训练集长什么样子的

imgs,labels=next(iter(train_dl))
imgs.shape

画图

import numpy as np
plt.figure(figsize=(20,5))
for i,img in enumerate(imgs[:20]):
    #进行轴变换
    npimg=img.numpy().transpose((1,2,0))
    plt.subplot(2,10,i+1)
    plt.imshow(npimg,cmap=plt.cm.binary)

    plt.axis('off')
plt.show()

继承torch的类，创建一个Model的类

import torch.nn.functional as F
num_classes=10
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        #特征提取网络
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
        self.pool1=nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.pool2=nn.MaxPool2d(2)
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3)
        self.pool3=nn.MaxPool2d(2)

        #分类网络
        self.fc1 = nn.Linear(512,256)
        self.fc2 = nn.Linear(256,num_classes)

        #前向传播
    def forward(self, x):
        x=self.pool1(self.conv1(x))
        x=self.pool2(self.conv2(x))
        x=self.pool3(self.conv3(x))

        x=F.relu(self.fc1(x))
        x=self.fc2(x)
        return x

采用模型训练

from torchinfo import summary
model = Model().to(device)
summary(model)

编写训练的函数

def train(dataloader, model, loss_fn, optimizer):
    size=len(dataloader.dataset)
    num_batches=len(dataloader)
    train_loss,train_acc=0,0
    for X,y in dataloader:
        X,y=X.to(device),y.to(device)

        pred=model(X)
        loss=loss_fn(pred,y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_acc+=(pred.argmax(1)==y).type(torch.float).sum().item()
        train_loss+=loss.item()
    train_acc/=size
    train_loss/=size
    return train_loss,train_acc

测试函数

def test(dataloader, model, loss_fn):
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    total_loss, correct = 0, 0

    with torch.no_grad():
        for imgs, labels in dataloader:
            imgs, labels = imgs.to(device), labels.to(device)

            pred = model(imgs)
            loss = loss_fn(pred, labels)

            total_loss += loss.item()
            correct += (pred.argmax(1) == labels).type(torch.float).sum().item()

    avg_loss = total_loss / num_batches
    accuracy = correct / size
    return avg_loss, accuracy

开始训练，训练5轮

# 训练循环
epochs = 5  # 增加训练轮数
train_loss = []
train_acc = []
test_loss = []
test_acc = []

best_acc = 0.0  # 保存最佳模型

for epoch in range(epochs):
    # 训练
    epoch_train_loss, epoch_train_acc = train(train_dl, model, loss_fn, optimizer)

    # 测试
    epoch_test_loss, epoch_test_acc = test(test_dl, model, loss_fn)

    # 更新学习率
    scheduler.step()

    # 记录结果
    train_loss.append(epoch_train_loss)
    train_acc.append(epoch_train_acc)
    test_loss.append(epoch_test_loss)
    test_acc.append(epoch_test_acc)

    # 保存最佳模型
    if epoch_test_acc > best_acc:
        best_acc = epoch_test_acc
        torch.save(model.state_dict(), 'best_model.pth')

    # 打印结果
    # if (epoch + 1) % 5 == 0 or epoch == 0:
    template = 'Epoch:{:3d}, LR:{:.4f}, Train Acc:{:.2f}%, Train Loss:{:.4f}, Test Acc:{:.2f}%, Test Loss:{:.4f}'
    current_lr = optimizer.param_groups[0]['lr']
    print(template.format(
        epoch + 1,
        current_lr,
        epoch_train_acc * 100,
        epoch_train_loss,
        epoch_test_acc * 100,
        epoch_test_loss
    ))

print(f'Finished Training. Best Test Accuracy: {best_acc * 100:.2f}%')

画图

from datetime import datetime
current_time = datetime.now()

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_acc, label='Train Accuracy')
plt.plot(train_acc, 'ro')
plt.plot(test_acc, label='Test Accuracy')
plt.plot(test_acc,'go')
plt.title('Accuracy')
plt.xlabel(current_time)
plt.grid(True)
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_loss, label='Train Loss')
plt.plot(train_loss, 'ro')
plt.plot(test_loss, label='Test Loss')
plt.plot(test_loss, 'go')
plt.title('Loss')
plt.xlabel(current_time)
plt.grid(True)
plt.legend()
plt.tight_layout()
# plt.savefig('training_curve.png')
plt.show()

运行截图