从零构建自定义数据集：PyTorch图像分类实战指南

在实际机器学习项目中我们经常遇到一个核心矛盾公开数据集虽然方便但往往与我们的具体业务场景存在偏差。无论是图像分类、文本分析还是语音识别直接使用现成数据集训练出的模型在真实业务数据上的表现常常不尽如人意。因此掌握从零开始构建、标注、管理并最终用于训练一个属于自己的数据集是每个希望将AI技术真正落地到业务中的开发者必须跨越的一道坎。这个过程远不止是收集一些文件那么简单它涉及数据采集、清洗、标注、格式转换、划分、增强以及最终与训练框架的集成每一步都隐藏着影响模型最终效果的细节。本文旨在为有一定机器学习基础但缺乏完整数据集构建经验的开发者提供一个清晰的实践指南。我们将以一个通用的图像分类任务为例贯穿数据生命周期的各个环节从最原始的数据收集开始到使用PyTorch框架加载自定义数据集完成训练验证。你将了解到如何设计一个可维护的数据目录结构如何使用工具进行高效标注如何处理常见的格式不兼容问题以及如何编写一个健壮的Dataset类来应对各种数据异常。最终你将获得一套可以复用于自己项目的、从数据到模型的可工作流程。1. 理解自定义数据集的完整生命周期与核心挑战在动手写代码之前我们需要先建立起对数据集构建全流程的宏观认知。一个高质量的数据集是模型成功的基石而构建过程本身就是一个微型的数据工程项目。1.1 数据集的生命周期阶段一个典型的自定义数据集构建与使用流程包含以下六个主要阶段它们环环相扣需求定义与规划明确你的模型要解决什么问题如“识别五种不同的工业零件”从而确定需要采集的数据类型图像、文本、音频、数据规模每类至少需要多少样本、以及标注的标准边界框、分类标签、分割掩码。数据采集与收集通过爬虫、设备拍摄、业务系统导出、公开数据筛选等方式获取原始数据。这个阶段数据的“脏乱差”是常态。数据清洗与预处理剔除无效数据损坏的图片、空文本、统一格式调整图像尺寸、统一音频采样率、进行初步的归一化或标准化。这是提升数据质量的关键一步。数据标注为清洗后的数据打上标签。这是最耗时但也最核心的环节标注的准确性和一致性直接决定模型性能的上限。数据集组织与划分将标注好的数据按照一定的目录结构存放并科学地划分为训练集、验证集和测试集以防止模型过拟合和评估偏差。数据集加载与迭代编写代码如PyTorch的Dataset和DataLoader将磁盘上的数据高效地读取到内存并进行必要的实时数据增强以供模型训练使用。1.2 构建过程中的常见挑战与应对思路数据不平衡某些类别的样本数量远多于其他类别。解决方案包括对多数类进行欠采样、对少数类进行过采样如复制、SMOTE算法、或在损失函数中引入类别权重。标注成本高人工标注费时费力。可以考虑主动学习让模型选择最需要标注的样本、半监督学习利用少量标注数据和大量未标注数据或寻找替代的弱监督信号。标注不一致不同标注员对同一数据的理解可能不同。必须制定详细、可操作的标注规范并进行多人交叉校验。数据泄露在划分数据集前进行了全局的预处理如使用整个数据集的均值和方差进行标准化导致测试集信息“泄露”到训练过程。必须确保预处理只在训练集上拟合参数然后应用到验证集和测试集。代码与数据耦合过紧数据路径、类别列表等硬编码在代码中导致数据稍作改动就需要修改代码。应通过配置文件或命令行参数来管理这些元信息。理解了这些挑战我们在后续的每个步骤中就会有意识地采取规避措施。2. 环境准备与项目结构设计工欲善其事必先利其器。一个清晰的项目结构能让后续所有工作井然有序避免混乱。2.1 基础环境与工具选择我们选择Python和PyTorch作为本次实践的技术栈因为它们在研究和工业界都有广泛的应用生态丰富。Python 3.8: 机器学习领域的主流语言。PyTorch 1.9: 选择其稳定版本它提供了torchvision等强大的工具库。标注工具对于图像分类轻量级的工具如labelImg目标检测、Label Studio通用或甚至一个简单的脚本配合文件夹分类即可。本文为了简化我们将采用“按文件夹分类”的方式这也是最直观的图像分类数据集组织方式。数据处理库PIL/Pillow图像处理、opencv-python、pandas管理标注文件。开发环境Jupyter Notebook可用于探索但最终建议使用PyCharm或VSCode等IDE进行项目管理。可以通过以下命令创建基础环境# 创建并激活虚拟环境以conda为例 conda create -n custom_dataset python3.9 conda activate custom_dataset # 安装核心依赖 pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu # 以CPU版本为例根据CUDA情况调整 pip install Pillow opencv-python pandas jupyter matplotlib scikit-learn2.2 设计可维护的项目目录结构一个糟糕的目录结构会让数据管理变成噩梦。推荐以下结构它清晰地分离了原始数据、处理后的数据、代码和输出。your_project/ ├── data/ # 所有数据相关文件 │ ├── raw/ # 原始收集的数据未经任何处理 │ │ ├── class_a/ # 原始A类图片 │ │ ├── class_b/ # 原始B类图片 │ │ └── .../ │ ├── processed/ # 清洗和预处理后的数据 │ │ ├── train/ # 训练集 │ │ │ ├── class_a/ │ │ │ ├── class_b/ │ │ │ └── .../ │ │ ├── val/ # 验证集 │ │ │ ├── class_a/ │ │ │ ├── class_b/ │ │ │ └── .../ │ │ └── test/ # 测试集可选也可从val划分 │ │ ├── class_a/ │ │ ├── class_b/ │ │ └── .../ │ └── annotations/ # 存放标注文件如CSV、JSON │ ├── train_labels.csv │ ├── val_labels.csv │ └── class_names.txt # 类别名称列表文件 ├── src/ # 源代码 │ ├── dataset.py # 自定义Dataset类 │ ├── train.py # 训练脚本 │ ├── utils.py # 工具函数如数据划分、可视化 │ └── config.py # 配置文件路径、超参数 ├── models/ # 保存训练好的模型权重 ├── logs/ # 训练日志、TensorBoard文件 ├── notebooks/ # Jupyter Notebook探索性分析 └── requirements.txt # 项目依赖列表这种结构的优势在于原始数据不变raw/目录下的数据永远保持原样所有处理都生成新文件到processed/便于追溯和重试。按用途隔离训练、验证、测试数据物理分离避免误操作。代码与数据分离src/目录只关心逻辑通过配置文件读取数据路径。3. 数据采集、清洗与标注实战假设我们的任务是构建一个“猫狗鸟类”三分类图像数据集。我们从网络收集了约1000张相关图片但尺寸不一混杂了一些非相关图片和损坏文件。3.1 数据清洗脚本示例首先我们需要清理data/raw/目录。编写一个清洗脚本src/clean_data.pyimport os from PIL import Image import shutil def clean_raw_data(raw_data_dir, processed_dir, min_size(50, 50), valid_extensions{.jpg, .jpeg, .png}): 清洗原始数据检查文件格式、图像可读性、最小尺寸并复制到处理目录。 Args: raw_data_dir: 原始数据根目录其子文件夹为类别名。 processed_dir: 处理后的数据输出根目录。 min_size: 图像最小允许的宽高。 valid_extensions: 允许的图片后缀集合。 for class_name in os.listdir(raw_data_dir): class_raw_path os.path.join(raw_data_dir, class_name) class_proc_path os.path.join(processed_dir, class_name) if not os.path.isdir(class_raw_path): continue # 跳过非目录文件 os.makedirs(class_proc_path, exist_okTrue) for filename in os.listdir(class_raw_path): filepath os.path.join(class_raw_path, filename) # 1. 检查文件后缀 _, ext os.path.splitext(filename) if ext.lower() not in valid_extensions: print(f跳过非图片文件: {filepath}) continue # 2. 尝试打开图像检查是否损坏 try: with Image.open(filepath) as img: img.verify() # 验证文件完整性 except (IOError, SyntaxError) as e: print(f损坏或无法读取的图像: {filepath}, 错误: {e}) continue # 3. 重新打开以检查尺寸 try: with Image.open(filepath) as img: if img.size[0] min_size[0] or img.size[1] min_size[1]: print(f图像尺寸过小 {img.size}: {filepath}) continue # 4. 复制合格文件到处理目录 shutil.copy2(filepath, os.path.join(class_proc_path, filename)) except Exception as e: print(f处理图像时发生意外错误: {filepath}, 错误: {e}) continue print(数据清洗完成。) if __name__ __main__: raw_root ../data/raw processed_root ../data/processed_all # 清洗后未划分的完整数据 clean_raw_data(raw_root, processed_root)运行此脚本后data/processed_all/目录下将只包含格式正确、未损坏且尺寸合格的图像并按类别文件夹组织。3.2 数据划分生成训练集、验证集和测试集接下来我们需要将清洗后的数据划分为训练集、验证集和测试集。务必在划分前进行随机打乱并考虑类别平衡。使用sklearn的train_test_split可以方便地实现。import os import random import shutil from sklearn.model_selection import train_test_split def split_dataset(data_all_dir, output_base_dir, train_ratio0.7, val_ratio0.15, test_ratio0.15, seed42): 将按类别文件夹组织的数据集划分为训练、验证、测试集。 Args: data_all_dir: 清洗后未划分的数据目录结构为 class_a/, class_b/... output_base_dir: 划分后数据的输出根目录其下会创建 train/, val/, test/。 train_ratio, val_ratio, test_ratio: 划分比例三者之和应为1。 seed: 随机种子保证划分可复现。 random.seed(seed) # 确保输出目录存在 train_dir os.path.join(output_base_dir, train) val_dir os.path.join(output_base_dir, val) test_dir os.path.join(output_base_dir, test) for d in [train_dir, val_dir, test_dir]: os.makedirs(d, exist_okTrue) for class_name in os.listdir(data_all_dir): class_dir os.path.join(data_all_dir, class_name) if not os.path.isdir(class_dir): continue # 获取当前类别的所有图片文件 images [f for f in os.listdir(class_dir) if os.path.isfile(os.path.join(class_dir, f))] if not images: continue # 划分 # 先分训练临时再从临时中分验证和测试 train_files, temp_files train_test_split(images, train_sizetrain_ratio, random_stateseed) # 计算验证集在剩余数据中的比例 val_ratio_adj val_ratio / (val_ratio test_ratio) val_files, test_files train_test_split(temp_files, train_sizeval_ratio_adj, random_stateseed) # 复制文件到对应目录 def copy_files(file_list, target_dir): target_class_dir os.path.join(target_dir, class_name) os.makedirs(target_class_dir, exist_okTrue) for f in file_list: src os.path.join(class_dir, f) dst os.path.join(target_class_dir, f) shutil.copy2(src, dst) copy_files(train_files, train_dir) copy_files(val_files, val_dir) copy_files(test_files, test_dir) print(f类别 {class_name} 划分完成: 训练{len(train_files)}张, 验证{len(val_files)}张, 测试{len(test_files)}张) if __name__ __main__: all_data_dir ../data/processed_all split_output_dir ../data/processed # 对应之前设计中的 processed/ 目录 split_dataset(all_data_dir, split_output_dir, train_ratio0.7, val_ratio0.15, test_ratio0.15)注意数据划分是模型评估可靠性的关键。测试集在整个训练和调参过程中应完全不可见仅用于最终评估。验证集用于训练过程中的模型选择和超参数调整。3.3 生成标注文件CSV格式虽然我们使用文件夹名作为标签但生成一个CSV标注文件有助于管理和后续可能的变化如使用更复杂的标注。同时创建一个class_names.txt文件记录类别顺序这对将字符串标签映射为整数ID至关重要。import os import csv import glob def create_annotation_csv(data_base_dir, output_csv_path, class_names_file): 为按 train/val/test 划分的数据创建CSV标注文件。 CSV格式split, filepath, label, label_index splits [train, val, test] # 首先从train目录获取所有类别并建立索引映射 train_dir os.path.join(data_base_dir, train) class_names sorted([d for d in os.listdir(train_dir) if os.path.isdir(os.path.join(train_dir, d))]) name_to_idx {name: idx for idx, name in enumerate(class_names)} # 保存类别名称文件 with open(class_names_file, w, encodingutf-8) as f: for name in class_names: f.write(name \n) print(f类别名称已保存至: {class_names_file}) # 创建CSV文件并写入数据 with open(output_csv_path, w, newline, encodingutf-8) as csvfile: writer csv.writer(csvfile) writer.writerow([split, relative_path, label, label_index]) # 表头 for split in splits: split_dir os.path.join(data_base_dir, split) if not os.path.exists(split_dir): continue for label in class_names: label_dir os.path.join(split_dir, label) if not os.path.isdir(label_dir): continue # 遍历该类别下的所有图片 for img_file in os.listdir(label_dir): if img_file.lower().endswith((.png, .jpg, .jpeg)): # 构建相对路径便于在不同环境中迁移 rel_path os.path.join(split, label, img_file) writer.writerow([split, rel_path, label, name_to_idx[label]]) print(f标注CSV文件已生成: {output_csv_path}) print(f共发现 {len(class_names)} 个类别: {class_names}) if __name__ __main__: data_dir ../data/processed csv_path ../data/annotations/dataset_annotations.csv class_names_path ../data/annotations/class_names.txt create_annotation_csv(data_dir, csv_path, class_names_path)运行后你会得到dataset_annotations.csv和class_names.txt。CSV文件内容示例如下split,relative_path,label,label_index train,train/cat/001.jpg,cat,0 train,train/dog/002.jpg,dog,1 val,val/bird/003.jpg,bird,2 ...至此一个结构清晰、标注明确的数据集已经准备就绪。4. 实现PyTorch自定义Dataset与DataLoader这是连接数据和模型训练的关键桥梁。我们需要创建一个继承自torch.utils.data.Dataset的类。4.1 自定义Dataset类详解在src/dataset.py中我们创建CustomImageDataset类。import torch from torch.utils.data import Dataset from PIL import Image import pandas as pd import os class CustomImageDataset(Dataset): 自定义图像分类数据集类从CSV标注文件读取数据。 def __init__(self, annotations_csv, img_dir, transformNone, target_transformNone): Args: annotations_csv (str): CSV标注文件路径。 img_dir (str): 图像根目录路径CSV中relative_path的基准路径。 transform (callable, optional): 应用于图像的变换/增强。 target_transform (callable, optional): 应用于标签的变换。 self.img_labels pd.read_csv(annotations_csv) self.img_dir img_dir self.transform transform self.target_transform target_transform # 可选缓存图像路径和标签到列表提高迭代速度对于大数据集需谨慎 # self.samples [(row[relative_path], row[label_index]) for _, row in self.img_labels.iterrows()] def __len__(self): return len(self.img_labels) def __getitem__(self, idx): 根据索引返回一个样本图像标签。 # 从DataFrame中获取一行数据 row self.img_labels.iloc[idx] img_path os.path.join(self.img_dir, row[relative_path]) label row[label_index] # 读取图像 try: image Image.open(img_path).convert(RGB) # 确保为三通道RGB except Exception as e: # 如果图像读取失败可以返回一个占位符图像或抛出异常 print(f无法读取图像 {img_path}: {e}) # 简单处理返回一个空白图像和-1标签需要在训练中处理 image Image.new(RGB, (224, 224), colorblack) label -1 # 应用变换 if self.transform: image self.transform(image) if self.target_transform: label self.target_transform(label) else: # 确保标签是Tensor label torch.tensor(label, dtypetorch.long) return image, label def get_class_distribution(self): 返回数据集中每个类别的样本数量用于检查平衡性。 return self.img_labels[label].value_counts().to_dict()关键点解释__init__中读取CSV文件保存路径和变换函数。使用Pandas读取CSV便于筛选例如后续可以轻松创建只包含训练集的数据集对象。__getitem__是核心方法根据索引idx定位到具体行拼接出完整的图像路径用PIL.Image.open读取。务必使用.convert(RGB)因为有些PNG图像可能是RGBA四通道与预训练模型输入不匹配。加入了基本的异常处理。在实际生产中图像读取失败是常见问题需要决定是跳过、记录还是用默认值替代。返回前应用transform数据增强和target_transform标签转换这里我们只是转为Tensor。4.2 定义数据变换Transform与增强Augmentation数据增强是提升模型泛化能力、防止过拟合的有效手段。我们使用torchvision.transforms。from torchvision import transforms # 定义训练和验证/测试阶段不同的变换策略 # 训练时增强 归一化 train_transform transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1), # 颜色抖动 transforms.ToTensor(), # 转换为Tensor并归一化像素值到[0,1] transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # ImageNet统计量使用预训练模型时常用 ]) # 验证/测试时仅进行必要的 resize、裁剪 和 归一化不做随机增强 val_test_transform transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪到224x224 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])为什么使用ImageNet的均值和标准差许多预训练的卷积神经网络如ResNet、VGG是在ImageNet数据集上训练的其输入数据经过了以此统计量进行的归一化。如果你从头开始训练或者你的数据分布与ImageNet差异极大可以计算自己数据集的均值和标准差。但对于微调Fine-tuning预训练模型使用其训练时的归一化参数通常是安全且方便的选择。4.3 创建DataLoader并可视化检查DataLoader负责从Dataset中按批次抽取数据并提供多进程加速、数据打乱等功能。from torch.utils.data import DataLoader # 假设我们已经有了CSV文件和根目录 train_csv ../data/annotations/train_labels.csv # 可以从总CSV中筛选出训练部分 val_csv ../data/annotations/val_labels.csv img_root_dir ../data/processed # 创建Dataset实例 train_dataset CustomImageDataset(annotations_csvtrain_csv, img_dirimg_root_dir, transformtrain_transform) val_dataset CustomImageDataset(annotations_csvval_csv, img_dirimg_root_dir, transformval_test_transform) # 创建DataLoader batch_size 32 num_workers 4 # 根据CPU核心数调整用于并行加载数据 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers, pin_memoryTrue) print(f训练集样本数: {len(train_dataset)}批次数量: {len(train_loader)}) print(f验证集样本数: {len(val_dataset)}批次数量: {len(val_loader)}) # 可视化检查一个批次的数据 import matplotlib.pyplot as plt import numpy as np def imshow(inp, titleNone): 从Tensor显示图像。 inp inp.numpy().transpose((1, 2, 0)) # 从(C, H, W)转为(H, W, C) mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) inp std * inp mean # 反归一化 inp np.clip(inp, 0, 1) plt.imshow(inp) if title: plt.title(title) plt.pause(0.001) # 暂停一下让绘图更新 # 获取一个批次 images, labels next(iter(train_loader)) print(f批次图像形状: {images.shape}) # 应为 [batch_size, 3, 224, 224] print(f批次标签形状: {labels.shape}) # 应为 [batch_size] # 显示网格图像 out torchvision.utils.make_grid(images) imshow(out, title[train_dataset.img_labels.iloc[i][label] for i in labels.numpy()])运行上述代码你应该能看到一个由猫、狗、鸟类图像组成的网格并且图像经过了随机裁剪、翻转等增强。这确认了数据管道工作正常。5. 构建训练循环与模型验证数据管道就绪后我们就可以将其接入标准的PyTorch训练流程。这里我们以一个简单的ResNet-18模型为例。5.1 定义模型、损失函数与优化器在src/train.py中import torch import torch.nn as nn import torch.optim as optim from torchvision import models from dataset import CustomImageDataset from torch.utils.data import DataLoader import time import copy def train_model(data_loaders, dataset_sizes, num_classes, device, num_epochs25): 训练和验证模型的通用函数。 Args: data_loaders (dict): 包含train和val两个DataLoader的字典。 dataset_sizes (dict): 包含train和val数据集大小的字典。 num_classes (int): 分类类别数。 device (torch.device): 训练设备CPU或GPU。 num_epochs (int): 训练轮数。 Returns: model: 训练好的模型。 val_acc_history: 验证集准确率历史记录。 # 1. 初始化模型使用预训练的ResNet18替换最后的全连接层 model models.resnet18(pretrainedTrue) # 加载预训练权重 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) # 替换为我们的分类头 model model.to(device) # 2. 定义损失函数和优化器 criterion nn.CrossEntropyLoss() # 只训练最后一层全连接层其他层学习率设小微调常用策略 optimizer optim.SGD([ {params: model.layer4.parameters(), lr: 1e-4}, # 深层参数小学习率 {params: model.fc.parameters(), lr: 1e-3} # 分类头大学习率 ], momentum0.9, weight_decay1e-4) # 学习率调度器每7个epoch衰减一次学习率 scheduler optim.lr_scheduler.StepLR(optimizer, step_size7, gamma0.1) # 3. 训练循环 since time.time() best_model_wts copy.deepcopy(model.state_dict()) best_acc 0.0 val_acc_history [] for epoch in range(num_epochs): print(fEpoch {epoch}/{num_epochs - 1}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, val]: if phase train: model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 迭代数据 for inputs, labels in data_loaders[phase]: inputs inputs.to(device) labels labels.to(device) # 清零梯度 optimizer.zero_grad() # 前向传播 # 只在训练阶段追踪历史以计算梯度 with torch.set_grad_enabled(phase train): outputs model(inputs) _, preds torch.max(outputs, 1) loss criterion(outputs, labels) # 反向传播 优化仅在训练阶段 if phase train: loss.backward() optimizer.step() # 统计 running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) if phase train: scheduler.step() # 更新学习率 epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 深拷贝模型如果验证准确率提升 if phase val and epoch_acc best_acc: best_acc epoch_acc best_model_wts copy.deepcopy(model.state_dict()) if phase val: val_acc_history.append(epoch_acc.cpu().numpy()) print() time_elapsed time.time() - since print(f训练完成于 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s) print(f最佳验证准确率: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, val_acc_history5.2 组织训练流程并启动在主函数中我们需要整合所有部分if __name__ __main__: import torchvision.transforms as transforms from torch.utils.data import DataLoader import os # 配置参数 data_dir ../data/processed annotation_dir ../data/annotations batch_size 32 num_workers 4 num_epochs 15 num_classes 3 # 猫、狗、鸟 # 设备配置 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 数据变换同上文 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), val: transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 创建Dataset和DataLoader image_datasets { x: CustomImageDataset( annotations_csvos.path.join(annotation_dir, f{x}_labels.csv), img_dirdata_dir, transformdata_transforms[x] ) for x in [train, val] } dataloaders { x: DataLoader(image_datasets[x], batch_sizebatch_size, shuffle(x train), num_workersnum_workers, pin_memoryTrue) for x in [train, val] } dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} # 获取类别名称可选用于打印 class_names image_datasets[train].img_labels[label].unique().tolist() print(f类别: {class_names}) # 训练模型 model, val_acc_history train_model( dataloaders, dataset_sizes, num_classes, device, num_epochsnum_epochs ) # 保存模型 torch.save(model.state_dict(), ../models/custom_dataset_resnet18.pth) print(模型已保存。) # 可以在这里添加测试集评估代码 # test_dataset CustomImageDataset(...) # test_loader DataLoader(...) # evaluate_model(model, test_loader, device)运行此脚本你将看到训练过程中每个epoch的训练和验证损失、准确率。模型权重会保存在models/目录下。6. 常见问题排查与最佳实践构建和训练自定义数据集时你会遇到各种问题。以下是典型问题及其排查路径。6.1 数据加载与预处理问题问题现象可能原因检查方式处理建议RuntimeError: stack expects each tensor to be equal size一个批次内的图像尺寸不一致。检查Dataset.__getitem__返回的图像是否都转换为Tensor且尺寸相同。确保在transform中包含了Resize或RandomResizedCrop等统一尺寸的操作。KeyError: ‘label_index’或 CSV 列名错误CSV文件列名与代码中读取的列名不匹配。打印self.img_labels.columns查看DataFrame的列名。修改__getitem__中的列名或重新生成CSV文件确保列名一致。图像显示为纯色或噪声归一化参数用错或图像数据范围不对。检查transforms.Normalize的mean和std参数并确认ToTensor()将图像转换到了[0,1]。使用transforms.ToPILImage()和反归一化代码可视化Tensor确认图像正常。内存占用过高或加载缓慢一次性将所有图像加载到内存num_workers设置过高。监控系统内存使用情况。Dataset应只存储文件路径在__getitem__中读取图像。调整DataLoader的num_workers通常为CPU核心数。使用pin_memoryTrue加速GPU传输。6.2 训练过程问题问题现象可能原因检查方式处理建议损失Loss不下降准确率随机波动学习率过高数据标签错误模型能力不足。检查初始损失值是否合理对于C类分类初始-log(1/C)。可视化几个批次的图像和标签。降低学习率如从1e-3降到1e-4。检查数据清洗和标注质量。尝试更复杂的模型。训练准确率很高但验证准确率很低过拟合模型过于复杂训练数据太少缺乏数据增强。对比训练集和验证集的损失/准确率曲线。增加数据增强的强度如随机旋转、裁剪、颜色抖动。添加Dropout层。使用更简单的模型。收集更多数据。验证准确率始终为0或极低验证集数据路径错误标签映射错误模型最后一层输出维度不对。单独运行验证集的一个批次打印预测结果和真实标签。检查model.fc.out_features是否等于num_classes。仔细核对验证集CSV文件和图像路径。检查class_names.txt与训练集是否一致。确保模型输出维度与类别数匹配。GPU内存溢出CUDA out of memory批次大小batch size过大模型或中间变量未释放。尝试减小batch_size。使用torch.cuda.empty_cache()。逐步减小batch_size如64-32-16。检查代码中是否有不必要的张量累积。使用梯度累积来模拟大批次。6.3 数据集构建最佳实践清单版本控制数据对原始数据raw/和处理脚本进行版本控制如Git LFS。processed/目录下的数据应由脚本生成不纳入版本控制。记录数据谱系在README.md或单独的文件中记录数据来源、清洗步骤、标注规则、划分比例和随机种子。这对于复现实验至关重要。始终保留测试集测试集只在最终评估模型时使用一次不要用它来调整模型或选择超参数。处理类别不平衡在构建DataLoader时可以使用WeightedRandomSampler为少数类样本赋予更高的采样权重或在损失函数nn.CrossEntropyLoss中设置weight参数。实现数据集缓存对于小数据集可以在__init__中将所有图像加载到内存或缓存到磁盘以加速训练。对于大数据集需确保__getitem__的IO效率。编写单元测试为你的Dataset类编写简单的测试确保它能正确返回指定索引的数据并且transform正常工作。分离配置与代码将数据路径、类别数、图像尺寸、批次大小等参数写入配置文件如config.py或YAML文件避免硬编码。7. 扩展方向与下一步成功构建并训练一个基础的自定义数据集模型后你可以从以下几个方向深化更复杂的标注类型本文使用的是图像分类。对于目标检测如YOLO、Faster R-CNN你需要使用labelImg等工具生成PASCAL VOC或COCO格式的XML/JSON标注文件并实现相应的Dataset类来读取边界框。非图像数据对于文本数据你需要处理分词、构建词表、序列填充等。PyTorch提供了torchtext库来简化这些流程。对于音频数据可能需要提取MFCC等特征。使用专业数据管理工具当数据量庞大或标注流程复杂时可以考虑使用DVC进行数据版本控制或使用Label Studio、CVAT等专业标注平台。自动化数据流水线将数据清洗、划分、增强等步骤封装成可复用的脚本或工作流如Apache Airflow DAG实现数据预处理自动化。探索高级增强技术使用albumentations库进行更丰富、更高效的图像增强或尝试基于深度学习的数据生成如GAN来扩充少数类样本。模型部署将训练好的模型使用TorchScript、ONNX格式导出并集成到Web服务如Flask、FastAPI或移动端应用中。构建自定义数据集是一个迭代的过程。第一版数据集和模型很少是完美的。重要的是建立起从数据到模型的可重复、可调试的完整管道然后在此基础上持续收集反馈、修正标注错误、补充困难样本从而驱动模型性能的不断提升。