摘要:
合集:AI案例-CV-动物
赛题:百度飞浆科技文化节 | 猫十二分类大赛
数据集:飞浆12种类型的猫图片数据集(2022)
解决方案:Pytorch框架、ResNet50模型
一、赛题简介
本项目来源于飞浆平台的图像分类学习赛。比赛要求参赛选手对十二种猫进行分类,属于CV方向经典的图像分类任务。
二、数据集内容
比赛主办方提供数据集包含12种猫的图片,并划分为训练集与测试集。
训练集: 提供高清彩色图片以及图片所属的分类,共有2,160张猫的图片,含标注文件。 测试集: 仅提供彩色图片,共有240张猫的图片,不含标注文件。
图片样例:

数据结构
数据样例:train_list.txt
cat_12_train/8GOkTtqw7E6IHZx4olYnhzvXLCiRsUfM.jpg 0
cat_12_train/hwQDH3VBabeFXISfjlWEmYicoyr6qK1p.jpg 0
cat_12_train/RDgZKvM6sp3Tx9dlqiLNEVJjmcfQ0zI4.jpg 0
cat_12_train/ArBRzHyphTxFS2be9XLaU58m34PudlEf.jpg 0
cat_12_train/kmW7GTX6uyM2A53NBZxibYRpQnIVatCH.jpg 0
...
cat_12_train/tXVja0Lc9gNOlepGZDJuwWKIMs3hPyTi.jpg 2
cat_12_train/dc5IFQ9MmiA0BpN4VnbTSrZlkPKaY27D.jpg 2
cat_12_train/7o5n6deIzLkFypcQjStBuTRfv2VAYiU8.jpg 2
cat_12_train/AH7XsrUvKOp3PGufhoIz9Fa01yqE24c5.jpg 2
cat_12_train/xqX8vwGdHkolSrcs7Ve2CaIQTgnFjZJB.jpg 2
数据集版权许可协议
三、ResNet50分类样例
安装开发包
参考文章《安装深度学习框架PyTorch》,选择合适的CUDA版本进行安装,例如:
conda create -n pytorch241-gpu python=3.10
conda activate pytorch241-gpu
conda install pytorch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 pytorch-cuda=12.1 -c pytorch -c nvidia
导入开发包
import random
import torch
import torch.nn as nn
import torchvision
from PIL import Image
from matplotlib import pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms as T, models
1、分割数据集
# 分割训练集
train_ratio = 0.7
train_paths, train_labels = [], []
valid_paths, valid_labels = [], []
with open('./train_list.txt', 'r') as f:
lines = f.readlines()
for line in lines:
if random.uniform(0, 1) < train_ratio:
train_paths.append(line.split(' ')[0])
label = line.split(' ')[1]
train_labels.append(int(line.split(' ')[1]))
else:
valid_paths.append(line.split(' ')[0])
valid_labels.append(int(line.split(' ')[1]))
2、定义TrainData
class TrainData(Dataset):
def __init__(self):
super().__init__()
self.color_jitter = T.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05)
self.normalize = T.Normalize(mean=0, std=1)
self.random_crop = T.RandomCrop(224, pad_if_needed=True)
self.random_rotation = T.RandomRotation(degrees=10)
def __getitem__(self, idx):
# 读取图片
image_path = train_paths[idx]
totensor = T.ToTensor()
image = Image.open(image_path)
image = image.convert('RGB')
image = totensor(image)
# 图像增广
features = self.color_jitter(image)
features = self.random_rotation(features)
features = self.random_crop(features)
features = self.normalize(features)
# 读取标签
labels = train_labels[idx]
return features, labels
def __len__(self):
return len(train_paths)
3、定义ValidData
class ValidData(Dataset):
def __init__(self):
super().__init__()
self.normalize = T.Normalize(mean=0, std=1)
def __getitem__(self, idx):
# 读取图片
image_path = valid_paths[idx]
totensor = T.ToTensor()
image = Image.open(image_path)
image = image.convert('RGB')
image = totensor(image)
# 图像变换
resize = T.Resize((256, 256))
features = resize(image)
features = self.normalize(features)
# 读取标签
labels = valid_labels[idx]
return features, labels
def __len__(self):
return len(valid_paths)
4、加载数据
train_dataloader =DataLoader(train_data, batch_size=32, shuffle=True)
valid_dataloader =DataLoader(valid_data, batch_size=16, shuffle=True)
train_data_size = len(train_data)
valid_data_size = len(valid_data)
5、定义模型
ResNet50 是一种基于残差学习(Residual Learning)框架的卷积神经网络(CNN),通常用于图像分类任务。它是 ResNet(Residual Network)家族的一部分,由微软研究院的 Kaiming He 等人在 2015 年提出。ResNet 通过引入“残差块”(Residual Block)来解决深度神经网络训练中的梯度消失问题,从而使得网络可以训练得更深,达到更好的性能。ResNet50 使用了一个包含 50 层的深度网络,它的基本思想是引入“跳跃连接”(skip connections),允许网络的某些层跳过前面的层,直接将信息传递给更深的层。这种跳跃连接使得网络能够更容易地进行反向传播,从而解决深度网络中的梯度消失问题。
这段代码完成了以下任务:
- 创建一个ResNet50模型实例。
- 修改模型的输出层以适应新的分类任务。
- 将模型和损失函数移动到GPU上(如果有GPU)。
- 初始化优化器以进行模型训练。
# 调用resnet并修改输出层神经元
model = models.resnet50(pretrained=False) # 不加载预训练权重
in_features = model.fc.in_features
num_classes = 12
new_fc = nn.Linear(in_features, num_classes)
model.fc = new_fc
model = model.cuda()
# loss and optimizer
# 定义损失函数
loss_function = nn.CrossEntropyLoss()
loss_f = loss_function.cuda()
# 初始化优化器
opt = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4)
6、训练数据
这段代码完成了以下任务:
- 训练模型20轮。
- 每轮训练过程中,计算并打印每个批次的损失。
- 每轮训练结束后,计算并打印该轮次的训练准确率。
- 每轮训练结束后,评估模型在验证集上的性能,计算并打印总损失和准确率。
- 训练结束后,保存训练好的模型参数。参考文章《模型和权重文件管理之PyTorch》。
# 开始训练
epoch = 20
for i in range(epoch):
model.train()
pe_train_acc = 0
for batch_id, data in enumerate(train_dataloader):
features, labels = data
features = features.cuda()
labels = labels.cuda()
predicts = model(features)
loss = loss_f(predicts, labels)
acc = (predicts.argmax(1) == labels).sum()
pe_train_acc = pe_train_acc + acc
loss.backward()
opt.step()
opt.zero_grad()
print('epoch: {}, batch: {}, loss: {}'.format(i + 1, batch_id + 1, loss.item()))
print("-------epoch: {}, acc: {}-------".format(i + 1, pe_train_acc/train_data_size))
model.eval()
print("--------eval--------")
acc = 0
total_test_loss = 0
total_right = 0
with torch.no_grad():
for data in valid_dataloader:
imgs, labels = data
imgs = imgs.cuda()
labels = labels.cuda()
predicts = model(imgs)
loss = loss_f(predicts, labels)
right = (predicts.argmax(1) == labels).sum().item()
total_test_loss = total_test_loss + loss.item()
total_right = total_right + right
total_acc = total_right/valid_data_size
print("total_loss: {}".format(total_test_loss))
print("total_acc: {}".format(total_acc))
torch.save(model.state_dict(), "resnet_to_cat12.pth")
训练过程
epoch: 1, batch: 1, loss: 2.4456729888916016
epoch: 1, batch: 2, loss: 2.6983845233917236
epoch: 1, batch: 3, loss: 2.4258038997650146
epoch: 1, batch: 4, loss: 2.571140766143799
epoch: 1, batch: 5, loss: 2.8005192279815674
epoch: 1, batch: 6, loss: 2.431760787963867
...
epoch: 1, batch: 46, loss: 2.388106107711792
epoch: 1, batch: 47, loss: 2.366446018218994
-------epoch: 1, acc: 0.12533512711524963-------
--------eval--------
total_loss: 119.31137657165527
total_acc: 0.07634730538922156
epoch: 2, batch: 1, loss: 2.3859503269195557
epoch: 2, batch: 2, loss: 2.234839916229248
...
7、推理
基于训练生成的模型进行推理。
源码开源协议
MIT License