《基于 PyTorch 的猫狗图像分类模型训练与评估》

2301_80883735

已于 2024-06-27 10:25:35 修改

阅读量1.1k

点赞数 16

文章标签： pytorch 分类人工智能

于 2024-06-22 13:56:56 首次发布

本文链接：https://blog.csdn.net/2301_80883735/article/details/139830894

版权

一、项目背景

随着图像数据的海量增长，对图像进行快速准确的分类成为了一项重要的任务。在本项目中，我们选择了常见的图像分类问题，例如区分猫和狗的图像，旨在通过深度学习技术，特别是基于卷积神经网络的方法，实现对这类图像的自动分类。

考虑到实际应用中的需求，我们需要一个能够在不同设备上运行，并且具有较高准确性和泛化能力的模型。通过使用流行的深度学习框架 PyTorch ，我们能够利用其丰富的功能和优化算法，来构建和训练我们的模型。

我们收集了大量的猫和狗的图像数据集，并对其进行了预处理和划分，分为训练集、验证集和测试集。通过精心设计的图像变换和数据增强策略，提高模型的鲁棒性和泛化能力

二、环境与库准备

首先，我们导入了一系列必要的库，包括 torch 相关的模块用于模型构建和训练，torchvision 用于数据处理和模型加载，以及一些数据处理和可视化的库。这些库为我们提供了构建、训练和评估模型所需的各种功能和工具。

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models
from torch.autograd import Variable
from tqdm import tqdm

import os
from PIL import Image
import cv2

三、设备选择

接下来，通过以下代码选择计算设备：

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DEVICE

这行代码会自动检测是否有可用的 GPU，如果有则使用 GPU（'cuda'）进行计算，否则使用 CPU。

四、数据预处理

定义了 transform111 字典，其中包含了针对训练数据和验证数据的不同预处理操作，如随机裁剪、旋转、翻转、标准化等，以增加数据的多样性和模型的泛化能力。

transform111 = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
        transforms.RandomRotation(degrees=15),
        transforms.RandomHorizontalFlip(),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(size=256),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}