PyTorch图像分类实战——基于ResNet18的RAF-DB情感识别（附完整代码和结果图）
关于作者
作者：小白熊
作者简介：精通python、matlab、c#语言，擅长机器学习，深度学习，机器视觉，目标检测，图像分类，姿态识别，语义分割，路径规划，智能优化算法，数据分析，各类创新融合等等。
联系邮箱：[email protected]
科研辅导、知识付费答疑、个性化定制以及其他合作需求请联系作者~
前言
在本文中，我们将详细介绍如何使用PyTorch框架，结合ResNet18模型，进行图像分类任务。这里我们选择了一个情感识别数据集——RAF-DB（Real-world Affective Faces Database），来进行实验。通过本文，你将学习到如何准备数据、构建模型、训练模型、评估模型，并可视化训练过程中的损失曲线。
1 模型理论
1.1 深度学习基础
深度学习是机器学习的一个分支，它通过使用深层神经网络来模拟人脑的学习过程。在图像分类任务中，卷积神经网络（Convolutional Neural Network, CNN）是最常用的模型之一。CNN通过卷积层、池化层、全连接层等结构，能够自动提取图像中的特征，并进行分类。
1.2 ResNet模型
ResNet（Residual Network）是一种深度卷积神经网络，它通过引入残差块（Residual Block）来解决深度神经网络中的梯度消失和梯度爆炸问题。残差块通过引入一个恒等映射（Identity Mapping），使得网络在训练过程中能够更容易地学习到特征。ResNet有多个版本，如ResNet18、ResNet34、ResNet50等，其中数字表示网络的层数。ResNet18作为ResNet系列中的一个轻量级模型，具备较好的性能和较低的计算复杂度，非常适合用于图像分类任务
1.3 交叉熵损失函数
在分类任务中，交叉熵损失函数（Cross Entropy Loss）是最常用的损失函数之一。它衡量的是模型输出的概率分布与真实标签的概率分布之间的差异。交叉熵损失函数越小，表示模型的预测结果越接近真实标签。
1.4 优化器
优化器用于更新模型的权重，以最小化损失函数。常用的优化器有SGD（随机梯度下降）、Adam等。SGD是最基础的优化器，它通过计算梯度来更新权重，但容易陷入局部最小值。Adam优化器结合了动量（Momentum）和RMSprop的思想，能够在训练过程中自适应地调整学习率，通常能够取得更好的效果。
2 代码解析
2.1 数据准备
首先，我们需要准备数据集。这里使用的是RAF-DB数据集，它是一个情感识别数据集，包含了多种情感标签的图像。
image_path = './data/RAF-DB' # 数据集路径，需修改
labels_num = 7 # 标签类别数量，需修改
我们使用datasets.ImageFolder来加载数据集，它会自动根据文件夹名称来划分标签。然后，我们定义了数据转换（data_transform），包括随机裁剪、随机水平翻转、归一化等操作。
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
2.2 模型构建
我们选择了ResNet18作为基模型，并修改了最后一层全连接层的输出维度，以适应我们的分类任务。
net = models.resnet18()
fc_input_feature = net.fc.in_features
net.fc = nn.Linear(fc_input_feature, labels_num)
然后，我们加载了预训练的权重，并删除了最后一层的权重，因为我们需要重新训练这一层。
pretrained_weight = torch.hub.load_state_dict_from_url(
url='https://download.pytorch.org/models/resnet18-5c106cde.pth', progress=True)
del pretrained_weight['fc.weight']
del pretrained_weight['fc.bias']
net.load_state_dict(pretrained_weight, strict=False)
2.3 训练过程
在训练过程中，我们使用了交叉熵损失函数和SGD优化器。同时，我们还设置了学习率调度器（scheduler），它会在每10个epoch后，将学习率乘以0.1。
criterion = nn.CrossEntropyLoss()
LR = 0.01
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
训练过程包括前向传播、计算损失、反向传播和更新权重。我们还使用了tqdm库来显示训练进度。
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader, file=sys.stdout)
acc1 = 0
for step, data in enumerate(train_bar):
images, labels = data
images = images.to(device)
labels = labels.to(device)
output = net(images)
optimizer.zero_grad()
loss = criterion(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
predicted = torch.max(output, dim=1)[1]
acc1 += torch.eq(predicted, labels).sum().item()
train_accurate = acc1 / train_num
# validate
net.eval()
with torch.no_grad():
val_bar = tqdm(validate_loader, file=sys.stdout)
val_loss = 0.0
acc = 0
for val_data in val_bar:
val_images, val_labels = val_data
val_images = val_images.to(device)
val_labels = val_labels.to(device)
output = net(val_images)
loss = criterion(output, val_labels)
val_loss += loss.item()
predict_y = torch.max(output, dim=1)[1]
acc += torch.eq(predict_y, val_labels).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f train_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate, train_accurate))
train_losses.append(running_loss / train_steps)
if val_accurate > best_acc:
best_acc = val_accurate
save_path_epoch = os.path.join(save_path, f"resnet_{
epoch + 1}_{
val_accurate}.pth")
torch.save(net.state_dict(), save_path_epoch)
2.4 损失曲线绘制
最后，我们绘制了训练过程中的损失曲线，以便观察模型的训练效果。
plt.figure(figsize=(10, 8))
plt.plot(range(1, epochs + 1), train_losses, label='train')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('resnet')
plt.legend()
plt.savefig('./runs/loss.jpg')
plt.show()
3 文件夹结构
为了使代码更加清晰和易于管理，建议使用以下文件夹结构：
project_root/
│
├── data/
│ └── RAF-DB/
│ ├── train/
│ │ ├── class1/
│ │ ├── class2/
│ │ ...
│ └── val/
│ ├── class1/
│ ├── class2/
│ ...
│
├── runs/ # 保存训练好的模型权重和损失曲线图
│ ├── loss.jpg
│ ├── resnet_1_xxx.pth
│ ...
│
├── train.py # 训练脚本
│
├── class_indices.json # 类别索引映射文件
│
├── requirements.txt # 项目依赖包
│
└── README.md # 项目说明文档
在这个结构中，data文件夹用于存放数据集，数据划分参考图像分类模型数据集划分教程：如何划分训练集和验证集
本文的class_indices.json类别索引映射文件如下：
{
"0": "anger",
"1": "disgust",
"2": "fear",
"3": "happiness",
"4": "neutral",
"5": "sadness",
"6": "surprise"
}
4 完整代码
import os
import sys
import json
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
from torchvision import transforms, datasets
import torch.optim as optim
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
import torchvision.models as models
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
image_path = './RAF-DB/RAF-DB' # 数据集路径，需修改
labels_num = 7 # 标签类别数量，需修改
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('./class_indices.json', 'w') as json_file: # json文件路径，需修改
json_file.write(json_str)
batch_size = 64 # 批处理大小，可修改
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=batch_size, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
# 初始化模型
net = models.resnet18()
fc_input_feature = net.fc.in_features
net.fc = nn.Linear(fc_input_feature, labels_num)
# load权重
pretrained_weight = torch.hub.load_state_dict_from_url(
url='https://download.pytorch.org/models/resnet18-5c106cde.pth', progress=True)
del pretrained_weight['fc.weight']
del pretrained_weight['fc.bias']
net.load_state_dict(pretrained_weight, strict=False)
net.to(device)
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
LR = 0.01
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
epochs = 100 # 训练轮数，可修改
best_acc = 0.0
save_path = './runs'
train_steps = len(train_loader)
train_losses = [] # 存储每个epoch的训练损失
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader, file=sys.stdout)
acc1 = 0
for step, data in enumerate(train_bar):
images, labels = data
images = images.to(device)
labels = labels.to(device)
output = net(images)
optimizer.zero_grad()
loss = criterion(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
predicted = torch.max(output, dim=1)[1]
acc1 += torch.eq(predicted, labels).sum().item()
train_accurate = acc1 / train_num
# validate
net.eval()
with torch.no_grad():
val_bar = tqdm(validate_loader, file=sys.stdout)
for val_data in val_bar:
val_images, val_labels = val_data
val_images = val_images.to(device)
val_labels = val_labels.to(device)
output = net(val_images)
loss = criterion(output, val_labels)
val_loss += loss.item()
predict_y = torch.max(output, dim=1)[1]
acc += torch.eq(predict_y, val_labels).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f train_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate, train_accurate))
train_losses.append(running_loss / train_steps)
if val_accurate > best_acc:
best_acc = val_accurate
save_path_epoch = os.path.join(save_path, f"resnet_{
epoch + 1}_{
val_accurate}.pth")
torch.save(net.state_dict(), save_path_epoch)
# 绘制损失曲线
plt.figure(figsize=(10, 8))
plt.plot(range(1, epochs + 1), train_losses, label='train')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('resnet')
plt.legend()
plt.savefig('./runs/loss.jpg')
plt.show()
print('Finished Training')
if __name__ == '__main__':
main()
5 结语
在本文中，我们深入探讨了如何使用PyTorch框架和ResNet18模型，结合RAF-DB数据集来实现图像情感识别。通过系统的数据预处理、模型构建、训练与优化，以及评估等步骤，我们成功地训练出了一个能够识别图像中人物情感的模型。
RAF-DB数据集作为本文的核心数据资源，展现出了其独特的价值。它是一个大规模、高质量的面部表情数据库，包含了数千张精挑细选的高分辨率面部图像，每张图像都配备了精确的表情标签。这些图像覆盖了高兴、悲伤、愤怒、惊讶等多种基本及复合表情，为模型训练提供了丰富的数据支持。此外，RAF-DB数据集还具备真实性、完整性、易用性和研究驱动等特点，确保了研究的普适性和可靠性，为情感分析、人脸识别以及表情识别等领域的研究者提供了宝贵的资源。
在模型选择方面，我们采用了ResNet18模型。ResNet（Residual Network，残差网络）是一种由微软亚洲研究院提出的深度神经网络结构，其核心在于通过残差连接（residual connections）解决了深层网络训练中的梯度消失和梯度爆炸问题。在本文中，我们利用ResNet18模型对RAF-DB数据集进行了训练，并成功地构建了一个情感识别模型。