图像分类任务-猫狗大战（计算机视觉领域区分猫和狗照片）

相关：

https://blog.csdn.net/weixin_74009895/article/details/143255069

代码：

#第一部分：导包
# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms, models
 
import numpy as np
import matplotlib.pyplot as plt
from tqdm import *
from torchinfo import summary
import os
import sys
from PIL import Image  # pip install Pillow
#第二部分:加载数据集
# 原始数据路径
data_path = "data/dogs-vs-cats/"
 
# 指定训练数据集的路径并获取该路径下所有文件的文件名列表
train_folder = data_path + 'train/'
train_filelist = os.listdir(train_folder)
 
# 输出训练数据集的文件数量和前20个文件名
print("train: {0} {1} \n".format(len(train_filelist), train_filelist[:20]))
 
# 指定测试数据集的路径并获取该路径下所有文件的文件名列表
test_folder = data_path + 'test/'
test_filelist = os.listdir(test_folder)
 
# 输出测试数据集的文件数量和前20个文件名
print("test: {0} {1} \n".format(len(test_filelist), test_filelist[:20]))
 
 
#第三部分；自定义数据集
# 自定义数据集
class DogsVSCats(Dataset):
    def __init__(self, data_dir, split="train", transform=None):
        # 拼接数据集中所有图片路径并计算总数
        imgs = [os.path.join(data_dir, img) for img in os.listdir(data_dir)]
        imgs_num = len(imgs)
 
        # 根据数据集划分（train/val/test），选择不同数据
        # 训练集：验证集设为8:2
        if split == 'train':
            self.image_paths = imgs[:int(0.8 * imgs_num)]
        elif split == 'val':
            self.image_paths = imgs[int(0.8 * imgs_num):]
        else:
            self.image_paths = imgs
        self.split = split
 
        # 根据数据集划分（train/val/test）定义默认数据转换方法
        if transform is None:
            if split == 'train':
                self.transform = transforms.Compose([
                    transforms.RandomHorizontalFlip(),    # 随机水平翻转
                    transforms.RandomRotation(10),        # 随机旋转
                    transforms.Resize((256, 256)),        # 调整图像大小
                    transforms.RandomCrop(224),           # 从随机位置裁剪指定大小
                    transforms.ToTensor(),                # 将数据转换为张量
                    # 对三通道数据进行归一化（均值，标准差），数值是从 ImageNet 数据集的百万张图片中随机抽样计算得到
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                ])
            else:
                self.transform = transforms.Compose([
                    transforms.Resize((224, 224)),        # 调整图像大小
                    transforms.ToTensor(),                # 将数据转换为张量
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
                ])
        else:
            self.transform = transform
 
    def __len__(self):
        return len(self.image_paths)
 
    def __getitem__(self, idx):
        # 加载图像数据并对图像数据进行转换
        image = Image.open(self.image_paths[idx]).convert('RGB')
        image = self.transform(image)
 
        # 从图片路径中提取标签信息并转换为张量数据
        filename = self.image_paths[idx].split("/")[-1]
        if self.split == 'test':
            label = int(filename.split('.')[0])
        else:
            label = 1 if 'dog' in filename else 0
        label = torch.tensor(label, dtype=torch.long)
 
        return image, label
#第四部分：验证数据集是否可用
dataset = DogsVSCats(
    'data/dogs-vs-cats/train/',
    transform=transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
)
#第五部分：定义一个根据数据集的绘图函数
# 定义绘图函数，传入dataset即可
def plot(dataset, shuffle=False, cmap=None):
    # 创建数据加载器
    dataloader = DataLoader(dataset, batch_size=16, shuffle=shuffle)
 
    # 取出一组数据
    images, labels = next(iter(dataloader))
 
    # 将通道维度(C)移动到最后一个维度，方便使用matplotlib绘图
    images = np.transpose(images, (0, 2, 3, 1))
 
    # 创建2x8的子图对象
    fig, axes = plt.subplots(nrows=4, ncols=4, figsize=(8, 8))
 
    # 遍历每个子图，绘制图像并添加子图标题
    for i, ax in enumerate(axes.flat):
        ax.imshow(images[i], cmap=cmap)
        ax.axis('off')  # 隐藏坐标轴
 
        if hasattr(dataset, 'classes'):  # 如果数据集有预定义的类别名称，使用该名称作为子图标题
            ax.set_title(dataset.classes[labels[i]], fontsize=12)
        else:  # 否则使用类别索引作为子图标题
            ax.set_title(labels[i], fontsize=12)
 
    plt.show()
plot(dataset)
 
#第六部分：分割数据集
# 依次定义训练集、验证集、测试集
train_dataset = DogsVSCats('data/dogs-vs-cats/train/', split='train')
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
 
val_dataset = DogsVSCats('data/dogs-vs-cats/train/', split='val')
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)
 
test_dataset = DogsVSCats('data/dogs-vs-cats/test/', split='test')
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=4)
 
# 查看数据集样本数量
print(len(train_dataset))
print(len(val_dataset))
print(len(test_dataset))
 
 
#第七部分：定义 ResNet50 模型
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
 
# 获取全连接层的输入特征数
num_features = model.fc.in_features
# 将全连接层替换成一个新的全连接层
model.fc = nn.Linear(in_features=num_features, out_features=2, bias=True)
 
#第八部分：模型训练
class Trainer:
    def __init__(self, model, train_loader, val_loader):
        # 初始化训练数据集和验证数据集的dataloader
        self.train_loader = train_loader
        self.val_loader = val_loader
 
        # 判断用的设备是 CPU 还是 GPU，并将模型移动到对应的计算资源设备上
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device)
 
        # 定义优化器、损失函数学习率调整器
        self.optimizer = optim.SGD(self.model.parameters(), lr=0.001)
        self.criterion = nn.CrossEntropyLoss()
        self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.95)
 
        # 记录训练过程中损失和验证过程中的准确率
        self.train_losses = []
        self.val_accuracy = []
 
    def train(self, num_epochs):
        # tqdm用于显示进度条并评估任务的时间开销
        for epoch in tqdm(range(num_epochs), file=sys.stdout):
            # 记录损失值
            total_loss = 0
 
            # 批量训练
            self.model.train()
            for images, labels in self.train_loader:
                # 将数据转移到指定计算资源设备上
                images = images.to(self.device)
                labels = labels.to(self.device)
 
                # 预测、损失函数、反向传播
                self.optimizer.zero_grad()
                outputs = self.model(images)
                loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()
 
                # 记录训练集loss
                total_loss += loss.item()
 
            # 更新优化器的学习率
            self.scheduler.step()
 
            # 计算验证集的准确率
            accuracy = self.validate()
 
            # 记录训练集损失和验证集准确率
            self.train_losses.append(np.log10(total_loss))  # 由于数值有时较大，这里取对数
            self.val_accuracy.append(accuracy)
 
            # 打印中间值
            if epoch % 5 == 0:
                tqdm.write("Epoch: {0} Loss: {1} Acc: {2}".format(
                    epoch, self.train_losses[-1], self.val_accuracy[-1]))
 
    def validate(self):
        # 测试模型，不计算梯度
        self.model.eval()
        total = 0
        correct = 0
 
        with torch.no_grad():
            for images, labels in self.val_loader:
                # 将数据转移到指定计算资源设备上
                images = images.to(self.device)
                labels = labels.to(self.device)
 
                # 预测
                outputs = self.model(images)
                # 记录验证集总数和预测正确数量
                total += labels.size(0)
                correct += (outputs.argmax(1) == labels).sum().item()
 
        # 返回准确率
        accuracy = correct / total
        return accuracy
 
# 创建一个 Trainer 类的实例
trainer = Trainer(model, train_loader, val_loader)
# 训练模型，迭代 20 个周期
trainer.train(num_epochs=20)
#第九部分：输出结果并绘图
# 最终准确率
print("准确率为：",trainer.val_accuracy[-1])
# 使用Matplotlib绘制损失和准确率的曲线图
plt.plot(trainer.train_losses, label='loss')
plt.plot(trainer.val_accuracy, label='accuracy')
plt.legend()
plt.show()
#第十部分：模型预测
def inference(trainer, test_loader):
    # 设置模型为评估模式
    trainer.model.eval()
 
    # 记录测试样本的id和模型预测结果
    ids = []
    predictions = []
 
    with torch.no_grad():
        # 遍历测试集，这里的labels其实是id
        for images, labels in test_loader:
            # 将数据转移到指定计算资源设备上
            images = images.to(trainer.device)
            # 对图片进行预测
            outputs = trainer.model(images)
            # 获取预测结果，torch.max函数返回值中的最大值项
            _, predicted = torch.max(outputs.data, 1)
            # 将测试样本的 ID 添加到列表中
            ids.extend(labels.numpy().tolist())
            # 将预测结果添加到列表中
            predictions.extend(predicted.cpu().numpy().tolist())
 
    return ids, predictions
# 对测试集进行预测，获得预测结果
ids, predictions = inference(trainer, test_loader)
 
# 输出前 20 个样本的预测结果
predictions[:20]
# [0, 1, 0, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0]
 
# 在分类正确的情况下，0.995和1.0的评估差异不大
print(F.binary_cross_entropy(torch.tensor([1.0]), torch.tensor([1.0])))
print(F.binary_cross_entropy(torch.tensor([0.995]), torch.tensor([1.0])))
 
#第十一部分：保存训练好的模型
# 保存整个模型
torch.save(model, 'resnet50_model.pth')