目录

引言

一、项目背景与数据集介绍

1.1 项目背景

1.2 数据集结构(结合用户提供的5张图片)

1.2.1 整体目录结构(图1、图2)

1.2.2 具体类别与图像(图3、图4、图5)

二、完整代码实现(附逐行注释)

2.1 环境准备与依赖导入

2.2 自动生成train.txt和test.txt文件

代码解析:

2.3 自定义数据集类(继承Dataset)

代码解析:

2.4 数据预处理与加载(DataLoader)

代码解析:

2.5 模型构建(CNN网络定义)

代码解析:

2.6 模型训练与测试

代码解析:

2.7 模型初始化与训练执行

代码解析:

三、代码分段深度解析

3.1 数据准备:为什么选择“按类别分文件夹”?

3.2 自定义Dataset类的核心方法

3.3 CNN模型的设计逻辑

3.4 训练过程中的关键细节

四、实验结果与优化建议

4.1 预期结果

4.2 常见问题与优化建议

五、总结


引言

在计算机视觉领域,图像分类是核心技术之一,其应用场景覆盖医疗影像诊断、安防监控、智能推荐等多个领域。卷积神经网络(Convolutional Neural Network, CNN)凭借其局部感知和权值共享的特性,成为图像分类任务的“标配”模型。本文将以​​本地自定义食物图像数据集​​为对象,手把手演示如何使用PyTorch框架搭建卷积神经网络,完成从数据准备、模型构建到训练测试的全流程实战。

本文的所有代码和数据集结构均基于用户提供的实际项目素材(如图1-5所示),包含详细的注释和分步解析,适合PyTorch初学者及计算机视觉爱好者参考。


一、项目背景与数据集介绍

1.1 项目背景

食物分类是图像分类任务的垂直细分场景,其核心目标是通过模型识别图像中的食物种类(如薯条、八宝粥、骨肉相连等)。与通用图像分类(如ImageNet)不同,食物分类的数据集通常规模较小但类别更聚焦(本案例包含20类食物),因此需要更精细的数据组织和模型调优。

1.2 数据集结构(结合用户提供的5张图片)

用户提供的项目数据集结构清晰,完全符合PyTorch的Dataset加载规范。通过图1-5的文件管理器界面,我们可以直观看到以下层级关系:

1.2.1 整体目录结构(图1、图2)
  • ​根路径​​:2、卷积神经网络/data/食物分类/food_dataset(图1、图2)
  • ​子集文件夹​​:train(训练集)和test(测试集),分别存储训练和验证用的图像数据(图2、图3)。
  • ​类别子文件夹​​:traintest下各有多个以食物命名的子文件夹(如“八宝粥”“巴旦木”“薯条”等)(图3),每个子文件夹内存储对应类别的JPEG图像(图4、图5)。

这种“train/test目录下按类别分子目录”的结构是PyTorch官方推荐的​​自定义数据集标准结构​​,便于通过代码自动读取图像路径和标签。

 

                           图1                                                                   图2

                         图3                                                                     图4

                                                               图·5

1.2.2 具体类别与图像(图3、图4、图5)
  • ​类别示例​​:train文件夹下包含“八宝粥”“巴旦木”“白萝卜”“骨肉相连”“薯条”等20类食物(图3),每类食物对应一个独立子文件夹,test测试集文件结构完全相同,文件夹最底部图片数量较训练集少很多。
  • ​图像示例​​:以“八宝粥”类别为例(图4、图5),其文件夹内包含6张JPEG格式的图像(文件名如img_八宝粥罐_22.jpeg),每张图像均为食物特写,尺寸可能不统一(需通过代码统一缩放)。

这种结构下,标签可以通过“子文件夹名称”自动生成(如“薯条”文件夹的索引为n,则该文件夹下所有图像的标签为n),无需额外标注文件(如CSV或JSON),极大简化了数据预处理流程。


二、完整代码实现(附逐行注释)

2.1 环境准备与依赖导入

首先需要安装必要的库,包括PyTorch、Pillow(图像处理)、matplotlib(可视化)等。本文假设已配置好PyTorch环境(支持CUDA或MPS加速)。

# 导入基础库:用于文件操作、数值计算等
import os
import numpy as np

# 导入PyTorch核心库及数据加载工具
import torch
from torch.utils.data import Dataset, DataLoader  # Dataset定义数据集,DataLoader批量加载数据

# 导入图像处理库:PIL用于打开、保存、显示图像
from PIL import Image

# 导入PyTorch的图像预处理工具(缩放、转Tensor等)
from torchvision import transforms

2.2 自动生成train.txt和test.txt文件

在PyTorch中,Dataset类通常需要读取一个包含图像路径和标签的文本文件(如train.txt)。本节通过遍历文件夹结构,自动生成这两个文件。

def generate_txt_files(root_dir, subset_dir, output_file):
    """
    生成训练集或测试集的路径-标签列表文件(.txt)
    
    参数:
        root_dir (str): 数据集根目录(如'./食物分类/food_dataset')
        subset_dir (str): 子集名称('train'或'test')
        output_file (str): 输出文件路径(如'./train.txt')
    """
    # 拼接子集的完整路径(如'./食物分类/food_dataset/train')
    subset_path = os.path.join(root_dir, subset_dir)
    # 打开输出文件(写入模式)
    with open(output_file, 'w', encoding='utf-8') as f:
        # 遍历子集路径下的所有目录和文件(os.walk递归遍历)
        for root, dirs, files in os.walk(subset_path):
            # 如果当前目录下有子目录(说明是类别文件夹的父级)
            if dirs:
                # 记录当前层级的所有类别名称(如['八宝粥', '巴旦木', ...])
                class_names = dirs
            else:
                # 当前目录是类别文件夹(无嵌套子目录),获取类别名称(父目录的最后一级)
                current_class = os.path.basename(root)
                # 计算当前类别对应的标签(类别名称在class_names中的索引)
                label = class_names.index(current_class)
                # 遍历当前类别文件夹下的所有图像文件
                for img_file in files:
                    # 拼接图像的完整路径(如'./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_22.jpeg')
                    img_path = os.path.join(root, img_file)
                    # 将路径和标签写入txt文件(格式:"路径 标签")
                    f.write(f"{img_path} {label}
")
    print(f"成功生成{subset_dir}集列表文件:{output_file}")

# 配置参数
dataset_root = r'.\食物分类\food_dataset'  # 数据集根目录(根据实际路径修改)
train_subset = 'train'                     # 训练集子集名称
test_subset = 'test'                       # 测试集子集名称

# 生成训练集和测试集的txt文件
generate_txt_files(dataset_root, train_subset, 'train.txt')
generate_txt_files(dataset_root, test_subset, 'test.txt')
代码解析:
  • ​函数功能​​:generate_txt_files通过os.walk递归遍历数据集目录,自动识别类别文件夹(如“八宝粥”),并将每个图像的路径与其对应的类别索引(标签)写入txt文件。
  • ​关键逻辑​​:
    • os.walk返回三元组(root, dirs, files),其中root是当前遍历的目录路径,dirs是当前目录下的子目录名列表,files是当前目录下的文件名列表。
    • dirs非空时,说明当前目录是类别文件夹的父级(如train),此时记录所有子目录名(类别名称);当dirs为空时,说明当前目录是类别文件夹(如train/八宝粥),通过os.path.basename(root)获取类别名称,并查找其在class_names中的索引作为标签。
    • 最终生成的train.txttest.txt格式示例如下:
      ./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_22.jpeg 0
      ./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_29.jpeg 0
      ./食物分类/food_dataset/train/巴旦木/img_巴旦木_01.jpeg 1
      ...

2.3 自定义数据集类(继承Dataset)

PyTorch的Dataset类是一个抽象类,需要重写__len__(数据集大小)和__getitem__(按索引获取数据)方法。本节定义FoodDataset类,用于加载图像和标签。

class FoodDataset(Dataset):
    def __init__(self, txt_path, transform=None):
        """
        初始化食物分类数据集
        
        参数:
            txt_path (str): 图像路径-标签列表文件的路径(如'train.txt')
            transform (callable, optional): 图像预处理变换(如缩放、转Tensor)
        """
        self.txt_path = txt_path
        self.transform = transform
        # 存储图像路径和标签的列表
        self.img_paths = []
        self.labels = []
        
        # 读取txt文件并解析数据
        with open(txt_path, 'r', encoding='utf-8') as f:
            # 逐行读取,每行格式为"图像路径 标签"
            for line in f.readlines():
                # 去除行首尾空白,按空格分割(注意:路径可能包含空格,需谨慎处理)
                parts = line.strip().split(' ')
                if len(parts) == 2:
                    img_path = parts[0]
                    label = int(parts[1])  # 标签转为整数
                    self.img_paths.append(img_path)
                    self.labels.append(label)
    
    def __len__(self):
        """返回数据集的大小(图像数量)"""
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        """
        按索引获取图像和标签
        
        参数:
            idx (int): 数据索引
            
        返回:
            tuple: (图像Tensor, 标签Tensor)
        """
        # 获取图像路径和标签
        img_path = self.img_paths[idx]
        label = self.labels[idx]
        
        # 用PIL打开图像(支持JPEG、PNG等格式)
        image = Image.open(img_path).convert('RGB')  # 确保图像为3通道(RGB)
        
        # 应用预处理变换(如缩放、转Tensor)
        if self.transform:
            image = self.transform(image)
        
        # 标签转为Tensor(PyTorch的损失函数需要Tensor类型)
        label = torch.tensor(label, dtype=torch.long)
        
        return image, label
代码解析:
  • ​初始化方法(__init__)​​:读取txt文件,将图像路径和标签分别存储到img_pathslabels列表中。注意处理路径中的空格问题(实际项目中建议避免路径包含空格)。
  • __len__方法​​:返回数据集的大小,即图像的总数,通过len(self.img_paths)获取。
  • __getitem__方法​​:
    • 根据索引idx获取对应的图像路径和标签。
    • PIL.Image.open打开图像,并通过.convert('RGB')确保图像为3通道(避免灰度图导致的通道数不一致问题)。
    • 应用预处理变换(如缩放、转Tensor),这部分在后续的data_transforms中定义。
    • 将标签转为long类型的Tensor(PyTorch的交叉熵损失函数CrossEntropyLoss要求标签为long类型)。

2.4 数据预处理与加载(DataLoader)

本节定义训练集和测试集的预处理变换,并通过DataLoader批量加载数据。

# 定义训练集和测试集的预处理变换
data_transforms = {
    'train': transforms.Compose([
        transforms.Resize([256, 256]),  # 调整图像大小为256x256(统一尺寸)
        transforms.ToTensor(),          # 转换为Tensor(自动归一化到[0,1])
    ]),
    'valid':
        transforms.Compose([
        transforms.Resize([256, 256]),
        transforms.ToTensor(),
    ]),
}#数组增强,

# 加载训练集和测试集
training_data = FoodDataset(
    file_path = './train.txt',          # 训练集的txt文件路径
    transform = data_transforms['train']  # 应用训练集的预处理
)
test_data = FoodDataset(
    file_path = './test.txt',           # 测试集的txt文件路径
    transform = data_transforms['valid']   # 应用测试集的预处理
)

# 创建数据加载器(DataLoader)
train_dataloader = DataLoader(training_data, batch_size=64,shuffle=True)#64张图片为一个包,
test_dataloader = DataLoader(test_data, batch_size=64,shuffle=True)
代码解析:
  • ​预处理变换(data_transforms)​​:
    • transforms.Resize([256, 256]):将图像尺寸统一为256x256,确保所有图像能组成批量Tensor(若图像尺寸不一致,无法进行批量运算)。
    • transforms.ToTensor():将PIL图像转换为PyTorch的Tensor,并自动将像素值从[0, 255]归一化到[0, 1]
  • ​DataLoader​​:
    • batch_size=64:每批加载64张图像,平衡内存占用和训练效率。
    • shuffle=True(训练集):打乱数据顺序,避免模型因数据顺序固定而学习到无关模式;测试集shuffle=False,便于按顺序评估结果。

2.5 模型构建(CNN网络定义)

本节定义一个简单的卷积神经网络(CNN),包含3个卷积层和2个全连接层,适用于食物分类任务。

import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):         # 输入大小 (1,3,256,256)(1,3,288,288)
        super(CNN, self).__init__()#
        self.conv1 = nn.Sequential(  #将多个层组合成一起。
            nn.Conv2d(          #2d一般用于图像,3d用于视频数据(多一个时间维度),1d一般用于结构化的序列数据
                in_channels=3,  # 图像通道个数,1表示灰度图(确定了卷积核 组中的个数),3表示RGB彩色图像
                out_channels=16,# 要得到几多少个特征图,卷积核的个数,
                kernel_size=5,  # 卷积核大小,5 * 5
                stride=1,       # 步长
                padding=2,      # 一般希望卷积核处理后的结果大小与处理前的数据大小相同,效果会比较好。那padding改如何设计呢?建议stride为1,kernel_size = 2*padding+1
            ),                  # 输出的特征图为 (1,16,256,256)(1,16,288,288)
            nn.ReLU(),  # relu层(1,16,256,256)(1,16,308,308)
            nn.MaxPool2d(kernel_size=2),  # 进行池化操作(2x2 区域), 输出结果为:(1,16,128,128)(1,16,144,144)
        )
        self.conv2 = nn.Sequential(  #输入
            nn.Conv2d(16, 32, 5, 1, 2),  # 输出 (1,32,128,128)          (1,32,144,144)
            nn.ReLU(),  # relu层(1,32,128,128)          (1,32,144,144)
            nn.Conv2d(32, 32, 5, 1, 2), # 输出(1,32,128,128)(1,32,144,144)
            nn.ReLU(),#输出(1,32,128,128)(1,32,144,144)
            nn.MaxPool2d(2),  # 输出             (1,32,64,64)              (1,32,72,72)
        )

        self.conv3 = nn.Sequential(  #输入        (32,64,64)
            nn.Conv2d(32, 128, 5, 1, 2),          #(1,128,64,64)          (1,128,72,72)
            nn.ReLU(),  # 输出(1,128,64,64)          (1,128,72,72)
        )

        self.out = nn.Linear(128 * 64 * 64,20)  # 全连接层得到的结果         (128 * 63 * 44,20)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)# 输出 (64,128, 64, 64)
        x = x.view(x.size(0), -1)  # flatten操作,结果为:(batch_size, 64 * 32 * 32)
        output = self.out(x)
        return output

model = CNN().to(device)
print(model)
代码解析:
  • ​卷积层设计​​:
    • nn.Conv2d:二维卷积层,用于提取图像的空间特征。in_channels为输入通道数(3表示RGB图像),out_channels为输出通道数(即卷积核数量,生成对应数量的特征图),kernel_size为卷积核尺寸(5x5),stride为步长(1表示每次滑动1个像素),padding为填充(2表示在图像边缘填充2个像素,确保输出尺寸与输入一致)。
    • nn.ReLU():修正线性单元激活函数,引入非线性特性,使模型能拟合更复杂的特征。
    • nn.MaxPool2d:最大池化层,通过取局部区域的最大值降低特征图尺寸(宽、高减半),减少计算量的同时保留主要特征。
  • ​全连接层​​:
    • nn.Linear(128 * 64 * 64,20):将卷积后的特征图展平为一维向量(长度为128 * 64 * 64),然后通过全连接层映射到20维的输出(对应20类食物的分类得分)。

2.6 模型训练与测试

本节定义训练函数和测试函数,完成模型的优化和性能评估。

def train(dataloader, model, loss_fn, optimizer):
    model.train()          #放开w参数的修改权限
#pytorch提供2种方式来切换训练和测试的模式,分别是:model.train() 和 model.eval()。
# 一般用法是:在训练开始之前写上model.trian(),在测试时写上 model.eval() 。
    batch_size_num = 1
    for X, y in dataloader:                 #其中batch为每一个数据的编号
        X, y = X.to(device), y.to(device)   #把训练数据集和标签传入cpu或GPU
        pred = model.forward(X)             #自动初始化 w权值
        loss = loss_fn(pred, y)             #通过交叉熵损失函数计算损失值loss
        # Backpropagation 进来一个batch的数据,计算一次梯度,更新一次网络
        optimizer.zero_grad()               #梯度值清零
        loss.backward()                     #反向传播计算得到每个参数的梯度值
        optimizer.step()                    #根据梯度更新网络参数

        loss = loss.item()                  #获取损失值
        if batch_size_num %1 == 0:
            print(f"loss: {loss:>7f}  [number:{batch_size_num}]")
        batch_size_num += 1

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()    #测试模式
    test_loss, correct = 0, 0
    with torch.no_grad():   #一个上下文管理器,关闭梯度计算。当你确认不会调用Tensor.backward()的时候。这可以减少计算所用内存消耗。
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model.forward(X)
            test_loss += loss_fn(pred, y).item() #
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            a = (pred.argmax(1) == y)  #dim=1表示每一行中的最大值对应的索引号,dim=0表示每一列中的最大值对应的索引号
            b = (pred.argmax(1) == y).type(torch.float)
    test_loss /= num_batches
    correct /= size
    print(f"Test result: \n Accuracy: {(100*correct)}%, Avg loss: {test_loss}")

loss_fn = nn.CrossEntropyLoss() #创建交叉熵损失函数对象,因为食物的类别是20,
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)#创建一个优化器,SGD为随机梯度下降算法??
# params:要训练的参数,一般我们传入的都是model.parameters()。
# lr:learning_rate学习率,也就是步长。


# train(train_dataloader, model, loss_fn, optimizer)
# test(test_dataloader, model, loss_fn)


epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    # test(test_dataloader, model, loss_fn)
print("Done!")
test(test_dataloader, model, loss_fn)
代码解析:
  • ​训练函数(train)​​:
    • model.train():开启训练模式,启用Dropout层(随机失活神经元)和BatchNorm层(计算当前批次的均值和方差)。
    • 遍历dataloader获取每个批次的图像(X)和标签(y),将数据和标签移动到目标设备(CPU/GPU)。
    • 前向传播计算预测值pred,通过损失函数loss_fn计算预测值与真实标签的损失loss
    • 反向传播loss.backward()计算梯度,优化器optimizer.step()更新模型参数。
    • 统计每个批次的损失值,按指定频率打印训练进度。
  • ​测试函数(test)​​:
    • model.eval():开启测试模式,禁用Dropout和BatchNorm的随机操作(使用训练阶段统计的均值和方差)。
    • with torch.no_grad():关闭自动梯度计算,减少内存消耗。
    • 遍历测试集,计算整体损失和准确率,评估模型的泛化能力。

2.7 模型初始化与训练执行

最后,初始化模型、损失函数和优化器,启动训练流程。

if __name__ == '__main__':
    # 配置参数
    epochs = 10            # 训练轮数
    # 自动选择计算设备(优先GPU,其次Apple MPS,最后CPU)
    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"Using {device} device")

    # 初始化模型、损失函数和优化器
    model = CNN().to(device)
    loss_fn = nn.CrossEntropyLoss()  # 交叉熵损失函数(适用于多分类)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam优化器
    
    # 启动训练
    print("开始训练...")
    for t in range(epochs):
        print(f"Epoch {t+1}\n-------------------------------")
        train(train_dataloader, model, loss_fn, optimizer)
    
    # 测试训练后的模型
    print("
训练结束,开始测试...")
    test(test_dataloader, model, loss_fn)
代码解析:
  • ​设备选择​​:通过条件判断自动检测可用的计算设备(GPU优先,其次Apple MPS,最后CPU),提升训练速度。
  • ​损失函数​​:nn.CrossEntropyLoss结合了Softmax激活函数和交叉熵损失,适用于多分类任务(输出为类别得分,无需手动计算概率)。
  • ​优化器​​:torch.optim.Adam是一种自适应学习率的优化器,通常比传统的SGD收敛更快。

三、代码分段深度解析

3.1 数据准备:为什么选择“按类别分文件夹”?

用户提供的图片(图1-5)展示了标准的“train/test目录下按类别分子目录”结构,这是PyTorch的Dataset加载机制的最佳实践,主要原因如下:

  • ​自动化标签生成​​:通过遍历子目录名(如“八宝粥”“薯条”),代码自动生成标签(子目录索引),无需手动标注。
  • ​兼容数据增强​​:在Dataset__getitem__方法中应用数据增强(如随机翻转),可以灵活地对每个批次的图像进行变换,避免重复存储增强后的图像。
  • ​符合官方最佳实践​​:PyTorch官方教程(如ImageFolder)推荐此结构,社区支持广泛,便于后续扩展(如使用预训练模型的ImageFolder加载方式)。

3.2 自定义Dataset类的核心方法

FoodDataset类继承自nn.ModuleDataset,其核心是__len____getitem__

  • __len__:告诉PyTorch数据集有多大,DataLoader通过此方法确定需要生成多少个批次。
  • __getitem__:定义如何获取单个样本。这里不仅返回图像和标签,还应用了预处理变换(如缩放、转Tensor),确保数据在输入模型前格式正确。

​注意​​:如果图像尺寸不一致,Resize变换是必须的(否则无法组成批量Tensor);如果数据量极大,可使用num_workers>0加速数据加载(但需注意Windows系统的多线程问题)。

3.3 CNN模型的设计逻辑

本文的CNN是一个典型的“卷积-池化-全连接”结构,设计逻辑如下:

  • ​浅层卷积(conv1)​​:提取边缘、纹理等低级特征(如食物的轮廓、颜色块)。
  • ​中层卷积(conv2)​​:组合低级特征形成中级特征(如食物的局部结构,如薯条的条状、八宝粥的颗粒感)。
  • ​深层卷积(conv3)​​:提取高级语义特征(如食物的整体形状、细节,如骨肉相连的骨头和肉的纹理)。
  • ​全连接层​​:将高级特征映射到类别空间,输出每个类别的得分。

​改进方向​​:实际项目中可使用更深的网络(如ResNet、VGG),或引入残差连接、注意力机制(如SE Block)提升性能。

3.4 训练过程中的关键细节

  • ​数据归一化​​:ToTensor变换将像素值从[0, 255]归一化到[0, 1]Normalize变换(若添加)进一步将数据标准化到均值为0、标准差为1的分布,有助于优化器更快收敛。
  • ​学习率调整​​:本文使用固定学习率(0.001),实际中可使用torch.optim.lr_scheduler动态调整(如余弦退火、阶梯衰减)。
  • ​过拟合处理​​:除了数据增强,还可添加Dropout层、L2正则化(weight_decay参数)等方法。

四、实验结果与优化建议

4.1 预期结果

假设数据集包含20类食物,每类约100张图像(训练集80%,测试集20%),使用上述代码训练10轮后,预期测试准确率可达70%-85%(具体取决于数据质量和类别复杂度)。

4.2 常见问题与优化建议

  • ​问题1:训练损失不下降​

    • 可能原因:学习率过大(模型无法收敛)或过小(收敛过慢)、数据预处理错误(如标签错误)、模型容量不足(网络太浅)。
    • 解决方法:调整学习率(如从0.001降至0.0001)、检查txt文件中的标签是否正确、增加卷积层或全连接层的神经元数量。
  • ​问题2:测试准确率远低于训练准确率​

    • 可能原因:过拟合(模型过度记忆训练数据)。
    • 解决方法:增加数据增强(如RandomHorizontalFlipRandomRotation)、添加Dropout层(如在conv1后加nn.Dropout2d(0.5))、使用早停法(Early Stopping)。
  • ​问题3:GPU内存不足​

    • 可能原因:批量大小(batch_size)过大、图像尺寸过大。
    • 解决方法:减小batch_size(如从64降至32)、降低图像尺寸(如从256x256降至128x128)、使用混合精度训练(torch.cuda.amp)。

五、总结

本文通过完整的代码示例和详细解析,演示了如何基于PyTorch框架和本地自定义数据集实现食物分类任务。核心步骤包括:

  1. 数据准备:按类别分文件夹组织数据,自动生成train.txttest.txt
  2. 自定义数据集:继承Dataset类,实现__len____getitem__方法。
  3. 模型构建:设计卷积神经网络,提取图像特征并映射到类别空间。
  4. 训练与测试:定义损失函数和优化器,通过DataLoader批量加载数据,迭代训练并评估模型性能。

通过本案例,读者可以掌握PyTorch处理自定义数据集的核心流程,并为后续的图像分类、目标检测等任务打下坚实基础。建议在实际项目中尝试不同的网络结构(如ResNet)、数据增强方法(如RandomCrop)和超参数调优(如学习率),以进一步提升模型性能。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐