CIFAR-10 图像分类

这是一个图像分类任务，主要是为了进一步熟悉pytorch训练流程，对自学进行一个检测。kaggle链接如下

CIFAR-10 - Object Recognition in Images | Kaggle

接下来对这个比赛进行分析

数据集

可以看到CIFAR-10是一个数据集Tiny Images Dataset的子集。Tiny Images Dataset这个数据集很大有八千万张图片。然后点击这个链接想了解下Tiny Images Dataset这个数据集。后面发现Tiny Images Dataset这个数据集永久下架，原因是数据集有偏见性。

CIFAR-10是一个具有60000张32×32彩色图片的数据集。该数据集有10大类，每个类有6000张图片，因此平衡性是不用担心的。整体有['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']这些类。整体有trainLabel.csv和test和train这些文件。

可以看到csv格式，给出的是除去后缀(png)后的图片名称。因此在进行数据预处理时，需要找到图片路径的话就要给出图片存放文件夹和后缀(png)。

数据预处理

这里我们的思路是使用torch.util.data.ImageFolder来加载图像数据。需要将训练数据train划分为training和val，按照8:2划分。划分后利用random.sample来随机化验证集索引，然后对train的图片遍历，若索引为val则目标目录path=os.path.join(train_path,(str)(id)+".png")。再将图片复制到相应目录。这样做的坏处是，同一时间磁盘会有两份数据，这对磁盘空间是一种浪费。但是这样处理比较简单，所以先这样使用。

import random
import shutil

import torch
import torch.nn as nn
import pandas as pd
import os
train_data=pd.read_csv('trainLabels.csv')
# print(train_data.shape)
images=train_data.iloc[:,0]
print(images)
os.makedirs('training',exist_ok=True)
os.makedirs('val',exist_ok=True)
split_ratio=0.2
train_path='train/train/'
test_path='test/test/'
num_sample=len(train_data)
num_train=int(num_sample*(1-split_ratio))
print("num of train:",num_train)
num_val=int(num_sample*split_ratio)
print("num of val:",num_val)
indices=random.sample(range(1,num_sample+1),num_val)
# print("indices:",min(indices))
for index,row in train_data.iterrows():
    print(index)
    id=row['id']
    path=os.path.join(train_path,(str)(id)+".png")
    label=row['label']
    if(id in indices):
        dir=os.path.join('val',label)
    else:
        dir=os.path.join('training',label)
    os.makedirs(dir,exist_ok=True)
    shutil.copy(path,dir)

处理好后结构如下

训练

def train(model, criterion, optimizer, train_loader, valid_loader, epochs, device):
    # 遍历每个epoch
    for epoch in range(epochs):
        model.train()  # 切换到训练模式
        running_loss = 0.0
        correct = 0
        total = 0
 #用 tqdm 库对数据加载器 train_loader 进行包装，生成一个带进度条的可迭代对象epoch_progress
        epoch_progress = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}")  # 批次进度条

        # 遍历训练集批次
        for batch_idx, (images, labels) in enumerate(epoch_progress):
            images, labels = images.to(device), labels.to(device)

            # 前向传播
            output = model(images)
            loss = criterion(output, labels)

            # 反向传播与优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 累计损失和准确率
            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(output.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            # 实时更新进度条信息（每10个批次更新一次，避免刷屏）
            if batch_idx % 10 == 0:
                batch_loss = loss.item()  # 当前批次的损失
                batch_acc = (predicted == labels).sum().item() / labels.size(0)  # 当前批次准确率
                epoch_progress.set_postfix({
                    "batch_loss": f"{batch_loss:.4f}",
                    "batch_acc": f"{batch_acc:.4f}",
                    "avg_loss": f"{running_loss / total:.4f}"  # 截至当前的平均损失
                })

        # 计算训练集整体指标
        train_loss = running_loss / len(train_loader.dataset)
        train_acc = correct / total

        # 在验证集上评估
        valid_loss, valid_acc = eval(model, criterion, valid_loader, device)

        # 打印当前epoch的总结信息
        print(f"\nEpoch {epoch + 1} 总结:")
        print(f"训练集 - 损失: {train_loss:.4f}, 准确率: {train_acc:.4f}")
        print(f"验证集 - 损失: {valid_loss:.4f}, 准确率: {valid_acc:.4f}\n")

    print('训练完成！')

这里使用tqdm 库对数据加载器 train_loader 进行包装，生成一个带进度条的可迭代对象epoch_progress，动态显示训练过程的准确率，损失。了解训练趋势。

评估

def eval(model, criterion, valid_loader, device):
    model.eval()  # 切换到评估模式
    correct = 0
    total = 0
    val_loss = 0.0

    # 显示验证过程进度条
    with torch.no_grad(), tqdm(valid_loader, desc="验证中") as val_progress:
        for images, labels in val_progress:
            images, labels = images.to(device), labels.to(device)
            output = model(images)
            loss = criterion(output, labels)

            # 累计验证损失和准确率
            val_loss += loss.item() * images.size(0)
            _, predicted = torch.max(output.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            # 实时更新验证进度条信息
            val_progress.set_postfix({
                "val_batch_loss": f"{loss.item():.4f}",
                "val_acc": f"{correct / total:.4f}"
            })

    # 计算验证集整体指标
    val_loss = val_loss / len(valid_loader.dataset)
    val_acc = correct / total
    return val_loss, val_acc

验证也是标注的图像分类验证过程，增加了进度条来观察训练过程。

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}\n")  # 显示当前训练设备

    # 定义数据转换
    traintrans = transforms.Compose([
        transforms.Resize(32),
        transforms.RandomHorizontalFlip(),  # 训练集数据增强
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    valtrans = transforms.Compose([
        transforms.Resize(32),
        transforms.CenterCrop(32),  # 验证集固定裁剪
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 加载数据集
    train_data = torchvision.datasets.ImageFolder(root='training', transform=traintrans)
    val_data = torchvision.datasets.ImageFolder(root='val', transform=valtrans)
    print(f"训练集样本数: {len(train_data)}, 验证集样本数: {len(val_data)}")
    print(f"类别数: {len(train_data.classes)}\n")  # 显示类别数量

    # 创建数据加载器
    batch_size = 256
    train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader = data.DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=2)

    # 定义模型、损失函数和优化器
    model = torchvision.models.resnet50(pretrained=False).to(device)
    model.fc = nn.Linear(in_features=model.fc.in_features, out_features=10).to(device)  # 10个类别
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 开始训练
    train(model, criterion, optimizer, train_loader, val_loader, epochs=300, device=device)

    # 保存模型
    torch.save(model.state_dict(), 'rs50_e100.pth')
    print("模型已保存为 'rs50_e100.pth'")

这里对用的是torchvision自带的resnet50,修改了最后一层全连接层的输出为10，对应CIFAR-10的10个类别

测试

import glob
import os
import re
import torch
import torch.nn as nn
import torchvision
from PIL import Image

#为了提取索引
train_transform = torchvision.transforms.ToTensor()
train_data = torchvision.datasets.ImageFolder(root='training',transform=train_transform)

class_to_idx = train_data.class_to_idx
idx_to_class = {v: k for k, v in class_to_idx.items()}

test_dir = 'test/test'
test_img_paths = glob.glob(os.path.join(test_dir, '*'))  # 获取所有图像路径（如.jpg/.png等）
def get_image_number(img_path):
    # 从路径中提取文件名（如"123.png"），再去掉扩展名".png"，转为整数
    filename = os.path.basename(img_path)
    return int(os.path.splitext(filename)[0]) #返回取消后缀的图像

# 按数字大小排序而不是字符串字典序
test_img_paths.sort(key=get_image_number)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model=torchvision.models.resnet50(pretrained=False).to(device)
model.fc=torch.nn.Linear(in_features=model.fc.in_features,out_features=10).to(device)
model.load_state_dict(torch.load('model.pth'))
model.eval()
transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
result=[]
with torch.no_grad():
    for img in test_img_paths:
        input=Image.open(img).convert('RGB')
        img_tensor = transform(input)  # 转换为Tensor
        img_tensor = img_tensor.unsqueeze(0)  # 增加批次维度 (1, C, H, W)
        img_tensor = img_tensor.to(device)
        output=model(img_tensor)
        print(output)
        _, predicted_idx = torch.max(output, 1)
        print(predicted_idx)
        predicted_class=idx_to_class[predicted_idx.item()]
        result.append(predicted_class)
with open("submission.csv",'w')as f:
    for img_path,label in zip(test_img_paths,result):
        img_num = get_image_number(img_path)
        f.write(f'{img_num},{label}\n')

这里添加了一个类别->索引，索引->类别的处理。因为模型输出的最后得到的是一个索引，根据这个索引再找到对应的类，这样才能找到具体类别。还需要注意，sort()默认是按照字符串字典序排序的，这样就会出现10.png在2.png前面这种情况，因此需要按照数字大小比较。