手把手教你用PyTorch做图像分类：5种服装识别，代码全中文注释

别再被“深度学习”吓跑啦这篇文章像教小朋友认卡片一样让电脑学会区分“上衣、鞋、包、下衣、手表”很多新手看到model.fit()、backward()这些术语就头疼。今天我不堆砌术语而是用最直白的语言 完整可运行的代码带你从头实现一个真正的图像分类器。你不需要是数学天才只要会一点Python基础就能跟着敲出一个能用的AI模型。我们最终要做成的事给电脑一张衣服/鞋/包的照片它能告诉你“这是上身衣服”、“这是鞋”还是“其他”。一、先别写代码——搞懂“卷积神经网络”在干嘛很多人学CNN死在第一步概念太多。我们换个比喻。1.1 把CNN想象成一个“找特征小组”卷积层Conv小组成员拿着不同的“模板”比如边缘模板、颜色模板在图片上滑来滑去记录每个位置匹配的程度。这些模板就是“卷积核”。池化层Pooling小组长说“每2x2的格子只保留最突出的那个其他扔掉”图片瞬间缩小一半但关键特征还在。全连接层FC最后所有特征汇总到一个“决策官”那里他根据各种特征的组合拍板决定“这是上衣”。1.2 我们的网络结构超级简单但有效为什么尺寸变成16x1664 → 卷积(padding1) → 64 → 池化 → 32 → 卷积 → 32 → 池化 → 16。16x16x16 4096这就是全连接层的输入大小。现在你心里有个地图了下面我们开始搭积木。二、准备工作项目文件夹结构建议你新建一个文件夹按下面结构创建文件和子文件夹my_fashion_classifier/ │ ├── common/ # 公共资源放数据和工具 │ ├── dataset/ # 所有图片放这里例如 0.jpg, 1.jpg, ... │ └── fashion-labels.csv # 标签文件告诉程序每张图属于哪类 │ ├── image_classification/ # 分类模块的所有代码 │ ├── __init__.py # 空文件标识这是一个Python包 │ ├── config.py # 配置参数就像游戏设置 │ ├── dataset.py # 读取图片和标签的代码 │ ├── model.py # CNN模型定义 │ ├── engine.py # 训练和验证的核心步骤 │ └── train.py # 主程序运行训练 │ └── classifier.pt # 训练完后会生成的模型文件 所有代码我都会给出完整版本你可以直接复制保存。数据集下载https://pan.baidu.com/s/1yvtFiQMfAMLumaqZGTtSbg?pwdvqmm三、配置文件集中管理所有“可调旋钮”创建 image_classification/config.py# ------------------- 路径配置 ------------------- IMAGE_DIR ../common/dataset/ # 图片文件夹路径 LABEL_CSV_PATH ../common/fashion-labels.csv # 标签CSV文件 # ------------------- 图像预处理 ------------------- IMAGE_SIZE 64 # 统一缩放成 64x64 像素 # ------------------- 数据划分 ------------------- TRAIN_RATIO 0.75 # 75%数据用于训练25%用于验证 RANDOM_SEED 42 # 随机种子保证结果可复现 # ------------------- 训练超参数 ------------------- BATCH_SIZE 32 # 每次喂给模型32张图显存不够可以改小比如16 EPOCHS 30 # 把所有数据完整看30遍 LEARNING_RATE 0.001 # 学习率每次调参的步长 # ------------------- 类别映射 ------------------- CLASS_NAMES { 0: 上身衣服, 1: 鞋, 2: 包, 3: 下身衣服, 4: 手表 } NUM_CLASSES len(CLASS_NAMES) # 自动计算类别数 5为什么把这些单独放因为后期调参比如想训练50轮、改学习率只需改这一个文件不用满代码翻找。四、读取数据让程序认识你的图片和标签4.1 标签文件长什么样fashion-labels.csv 有两列idtarget0011223344......id 对应图片文件名不含扩展名比如 0.jpg 对应 id0target 就是类别编号0上身衣服1鞋2包3下身衣服4手表4.2 写一个“数据集类”创建 image_classification/dataset.py内容如下关键地方都有注释import os import re import pandas as pd from PIL import Image from torch.utils.data import Dataset def natural_sort_key(filename): 让文件名按数字顺序排序而不是字符串顺序。 例如[1.jpg,2.jpg,10.jpg] 会变成 1,2,10 而不是 1,10,2 # 把文件名中的数字部分转换成整数非数字部分转小写 convert lambda text: int(text) if text.isdigit() else text.lower() # 用正则把数字和非数字拆开例如 img10.jpg - [img,10,.jpg] alphanum_key lambda key: [convert(c) for c in re.split(([0-9]), key)] return alphanum_key(filename) class FashionDataset(Dataset): 自定义数据集读取图片和对应的标签 def __init__(self, image_dir, label_csv_path, transformNone): 参数: image_dir (str): 存放图片的文件夹路径 label_csv_path (str): 标签CSV文件路径 transform (callable, optional): 对图片的预处理操作 self.image_dir image_dir self.transform transform # 获取文件夹下所有图片文件名并按数字顺序排序 self.image_names sorted(os.listdir(image_dir), keynatural_sort_key) # 读取标签CSV文件 labels_df pd.read_csv(label_csv_path) # 将id-target映射为字典方便快速查找 self.label_dict dict(zip(labels_df[id], labels_df[target])) def __len__(self): 返回数据集总共有多少张图片 return len(self.image_names) def __getitem__(self, idx): 根据索引 idx 返回一个样本 (image_tensor, label) # 1. 加载图片 img_path os.path.join(self.image_dir, self.image_names[idx]) image Image.open(img_path).convert(RGB) # 转为RGB三通道 # 2. 获取标签 # 注意CSV中的id与文件名索引是对应的假设id从0开始顺序排列 label self.label_dict[idx] # 3. 预处理缩放、转张量等 if self.transform is not None: image self.transform(image) else: raise RuntimeError(请提供transform预处理函数) return image, label4.3 为什么要用自然排序os.listdir() 返回的文件名顺序可能是 1.jpg, 10.jpg, 2.jpg。如果不排序索引0对应 1.jpg索引1对应 10.jpg但你的CSV里id1对应的是第二张图吗不是。所以必须按数字顺序排序让第0个文件名是 0.jpg如果图片命名从0开始这样索引才能和id对应上。如果你的图片命名不是从0开始的连续数字那么你需要根据文件名中的数字去匹配CSV中的id。但这里为了简化我们约定图片名和id一一对应且按顺序排列。4.4 创建DataLoader数据投喂器在主程序后面会写中我们会这样使用from torch.utils.data import DataLoader, random_split import torchvision.transforms as T from image_classification.dataset import FashionDataset from image_classification.config import * # 定义预处理缩放 转张量 transform T.Compose([ T.Resize((IMAGE_SIZE, IMAGE_SIZE)), T.ToTensor() # 将PIL图片转为 [0,1] 范围的张量形状 (C, H, W) ]) # 加载全部数据 full_dataset FashionDataset(IMAGE_DIR, LABEL_CSV_PATH, transformtransform) # 划分训练集(75%)和验证集(25%) train_size int(TRAIN_RATIO * len(full_dataset)) val_size len(full_dataset) - train_size train_dataset, val_dataset random_split(full_dataset, [train_size, val_size]) # 创建DataLoader train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, shuffleTrue, drop_lastTrue) val_loader DataLoader(val_dataset, batch_sizeBATCH_SIZE, shuffleFalse) # 测试一下取一个批次看看形状 for images, labels in train_loader: print(一个批次图片的形状:, images.shape) # 期望 [32, 3, 64, 64] print(一个批次标签的形状:, labels.shape) # 期望 [32] break输出示例一个批次图片的形状: torch.Size([32, 3, 64, 64]) 一个批次标签的形状: torch.Size([32])解释[32, 3, 64, 64] 32张图 × 3通道(RGB) × 高64像素 × 宽64像素。五、搭建CNN模型积木块详解创建 image_classification/model.pyimport torch.nn as nn import torch.nn.functional as F class FashionClassifier(nn.Module): 一个简单的CNN分类器 两个卷积层 两个池化层 一个全连接层 def __init__(self, num_classes5): super(FashionClassifier, self).__init__() # 第一个卷积层: 输入3通道(RGB) - 输出8个特征图 self.conv1 nn.Conv2d(in_channels3, out_channels8, kernel_size3, stride1, padding1) # 最大池化层: 2x2窗口, 步长2 - 尺寸减半 self.pool nn.MaxPool2d(kernel_size2, stride2) # 第二个卷积层: 输入8通道 - 输出16个特征图 self.conv2 nn.Conv2d(in_channels8, out_channels16, kernel_size3, stride1, padding1) # 全连接层: 输入特征数 16 * 16 * 16 4096, 输出 num_classes self.fc nn.Linear(16 * 16 * 16, num_classes) def forward(self, x): 前向传播: 输入 x 的形状 (batch_size, 3, 64, 64) # 第一块: 卷积 ReLU 池化 x F.relu(self.conv1(x)) # 形状: (batch, 8, 64, 64) x self.pool(x) # 形状: (batch, 8, 32, 32) # 第二块: 卷积 ReLU 池化 x F.relu(self.conv2(x)) # 形状: (batch, 16, 32, 32) x self.pool(x) # 形状: (batch, 16, 16, 16) # 展平: 把每个样本的特征图拉直成一维向量 x x.view(x.size(0), -1) # 形状: (batch, 4096) # 全连接层输出5个类别的原始分数(logits) x self.fc(x) # 形状: (batch, 5) # 输出 log_softmax (数值稳定的对数概率) # 注意: 训练时通常用 CrossEntropyLoss, 它内部已经包含 softmax, # 但这里为了演示, 我们直接输出 log_softmax 也可以配合 NLLLoss。 # 后面训练时我们使用 CrossEntropyLoss, 所以不需要在这里做 softmax。 return x # 返回原始 logits为什么全连接层的输入是 4096经过两次池化特征图尺寸从 64x64 → 32x32 → 16x16。通道数最后是16所以总特征数 16 × 16 × 16 4096。如果你不想手动计算可以用 nn.AdaptiveAvgPool2d((1,1)) 自适应全局池化那样全连接层输入就只是通道数16更简单。但这里为了让你理解尺寸变化我们保留了显式计算。六、训练引擎核心训练/验证循环创建 image_classification/engine.pyimport torch def train_one_epoch(model, train_loader, loss_fn, optimizer, device): 在训练集上训练一个epoch一轮 返回: 平均训练损失 model.train() # 设置为训练模式对Dropout/BatchNorm有影响 total_loss 0.0 num_batches 0 for images, labels in train_loader: # 把数据搬到GPU/CPU上 images images.to(device) labels labels.to(device) # 清空梯度否则梯度会累积 optimizer.zero_grad() # 前向传播模型预测 outputs model(images) # 计算损失 loss loss_fn(outputs, labels) # 反向传播计算梯度 loss.backward() # 更新模型参数 optimizer.step() total_loss loss.item() * images.size(0) # 累加该批次的总损失 num_batches 1 avg_loss total_loss / len(train_loader.dataset) # 平均每个样本的损失 return avg_loss def validate(model, val_loader, loss_fn, device): 在验证集上评估模型不更新参数 返回: 平均验证损失 model.eval() # 设置为评估模式 total_loss 0.0 num_batches 0 # 不计算梯度节省内存和计算 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels labels.to(device) outputs model(images) loss loss_fn(outputs, labels) total_loss loss.item() * images.size(0) num_batches 1 avg_loss total_loss / len(val_loader.dataset) return avg_loss为什么要区分 model.train() 和 model.eval()因为某些层如Dropout、BatchNorm在训练和推理时的行为不同。训练时需要随机丢弃神经元Dropout或计算批次统计量BatchNorm而推理时要用固定的统计量。七、主训练脚本把所有零件组装起来创建 image_classification/train.pyimport torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as T from torch.utils.data import DataLoader, random_split # 导入自定义模块 from config import * from dataset import FashionDataset from model import FashionClassifier from engine import train_one_epoch, validate def set_seed(seed): 设置随机种子保证结果可复现 torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def main(): # 1. 检查是否有GPU device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 2. 设置随机种子 set_seed(RANDOM_SEED) # 3. 数据预处理 transform T.Compose([ T.Resize((IMAGE_SIZE, IMAGE_SIZE)), T.ToTensor() ]) # 4. 加载数据集 print(正在加载数据集...) full_dataset FashionDataset(IMAGE_DIR, LABEL_CSV_PATH, transformtransform) train_size int(TRAIN_RATIO * len(full_dataset)) val_size len(full_dataset) - train_size train_dataset, val_dataset random_split(full_dataset, [train_size, val_size]) train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, shuffleTrue, drop_lastTrue) val_loader DataLoader(val_dataset, batch_sizeBATCH_SIZE, shuffleFalse) print(f训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}) # 5. 创建模型、损失函数、优化器 model FashionClassifier(num_classesNUM_CLASSES).to(device) loss_fn nn.CrossEntropyLoss() # 交叉熵损失内部包含softmax optimizer optim.AdamW(model.parameters(), lrLEARNING_RATE) # 6. 训练循环 best_val_loss float(inf) print(开始训练...) for epoch in range(1, EPOCHS 1): train_loss train_one_epoch(model, train_loader, loss_fn, optimizer, device) val_loss validate(model, val_loader, loss_fn, device) print(fEpoch {epoch:2d}/{EPOCHS} | Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f}) # 保存验证损失最低的模型 if val_loss best_val_loss: best_val_loss val_loss torch.save(model.state_dict(), classifier.pt) print(f - 保存新最佳模型 (val_loss{val_loss:.6f})) print(训练完成最佳模型已保存为 classifier.pt) if __name__ __main__: main()训练输出损失逐渐下降使用设备: cuda 正在加载数据集... 训练集大小: 18639, 验证集大小: 6214 开始训练... Epoch 1/30 | Train Loss: 0.213366 | Val Loss: 0.077724 - 保存新最佳模型 (val_loss0.077724) Epoch 2/30 | Train Loss: 0.076589 | Val Loss: 0.098107 Epoch 3/30 | Train Loss: 0.062420 | Val Loss: 0.049305 - 保存新最佳模型 (val_loss0.049305) Epoch 4/30 | Train Loss: 0.050954 | Val Loss: 0.047376 - 保存新最佳模型 (val_loss0.047376) Epoch 5/30 | Train Loss: 0.042681 | Val Loss: 0.043416 - 保存新最佳模型 (val_loss0.043416) Epoch 6/30 | Train Loss: 0.039183 | Val Loss: 0.048987 Epoch 7/30 | Train Loss: 0.033624 | Val Loss: 0.037801 - 保存新最佳模型 (val_loss0.037801) Epoch 8/30 | Train Loss: 0.029378 | Val Loss: 0.043007 Epoch 9/30 | Train Loss: 0.027791 | Val Loss: 0.039153 Epoch 10/30 | Train Loss: 0.023849 | Val Loss: 0.042653 Epoch 11/30 | Train Loss: 0.021031 | Val Loss: 0.040328 Epoch 12/30 | Train Loss: 0.019363 | Val Loss: 0.039726 Epoch 13/30 | Train Loss: 0.013487 | Val Loss: 0.061417 Epoch 14/30 | Train Loss: 0.014416 | Val Loss: 0.051996 Epoch 15/30 | Train Loss: 0.011395 | Val Loss: 0.046527 Epoch 16/30 | Train Loss: 0.012894 | Val Loss: 0.053243 Epoch 17/30 | Train Loss: 0.009627 | Val Loss: 0.049918 Epoch 18/30 | Train Loss: 0.011956 | Val Loss: 0.055744 Epoch 19/30 | Train Loss: 0.008906 | Val Loss: 0.047444 Epoch 20/30 | Train Loss: 0.007415 | Val Loss: 0.055333 Epoch 21/30 | Train Loss: 0.004162 | Val Loss: 0.052728 Epoch 22/30 | Train Loss: 0.007470 | Val Loss: 0.052018 Epoch 23/30 | Train Loss: 0.006289 | Val Loss: 0.058021 Epoch 24/30 | Train Loss: 0.007052 | Val Loss: 0.062240 Epoch 25/30 | Train Loss: 0.006248 | Val Loss: 0.054409 Epoch 26/30 | Train Loss: 0.004461 | Val Loss: 0.060665 Epoch 27/30 | Train Loss: 0.003012 | Val Loss: 0.068889 Epoch 28/30 | Train Loss: 0.007218 | Val Loss: 0.068075 Epoch 29/30 | Train Loss: 0.003673 | Val Loss: 0.053514 Epoch 30/30 | Train Loss: 0.007567 | Val Loss: 0.058273 训练完成最佳模型已保存为 classifier.pt注意实际损失数值取决于你的数据量和难度但趋势应该是不断下降并趋于平稳。八、用训练好的模型做预测训练完成后我们可以单独写一个脚本来测试模型效果。创建 image_classification/predict.pyimport torch import torchvision.transforms as T from PIL import Image from model import FashionClassifier from config import CLASS_NAMES, IMAGE_SIZE def predict_image(model, image_path, device): 对单张图片进行分类预测 # 预处理 transform T.Compose([ T.Resize((IMAGE_SIZE, IMAGE_SIZE)), T.ToTensor() ]) image Image.open(image_path).convert(RGB) input_tensor transform(image).unsqueeze(0) # 增加batch维度 - (1,3,64,64) input_tensor input_tensor.to(device) # 推理 model.eval() with torch.no_grad(): output model(input_tensor) # 输出 logits probabilities torch.softmax(output, dim1) # 转为概率 predicted_class torch.argmax(probabilities, dim1).item() return predicted_class, probabilities.cpu().numpy()[0] if __name__ __main__: # 加载训练好的模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model FashionClassifier(num_classeslen(CLASS_NAMES)) model.load_state_dict(torch.load(classifier.pt, map_locationdevice)) model.to(device) # 测试一张图片修改路径 test_image ../common/dataset/123.jpg # 替换成你的图片路径 class_id, probs predict_image(model, test_image, device) print(f预测类别: {CLASS_NAMES[class_id]}) print(各类别概率:) for i, name in CLASS_NAMES.items(): print(f {name}: {probs[i]:.4f})运行示例输出预测类别: 包 各类别概率: 上身衣服: 0.0000 鞋: 0.0000 包: 1.0000 下身衣服: 0.0000 手表: 0.0000完整代码下载包含数据集https://pan.baidu.com/s/15iajG1KBQr-39PYtG9_JiQ?pwdp8tw九、常见问题与解决方法QAQ1运行时报错FileNotFoundError: [Errno 2] No such file or directory: ../common/dataset/解决检查你的路径设置。建议使用绝对路径或在代码中动态获取当前文件所在目录。可以这样修改 config.pyimport os BASE_DIR os.path.dirname(os.path.dirname(os.path.abspath(__file__))) IMAGE_DIR os.path.join(BASE_DIR, common, dataset) LABEL_CSV_PATH os.path.join(BASE_DIR, common, fashion-labels.csv)Q2显存不足CUDA out of memory解决减小 BATCH_SIZE比如从32改成16或8。Q3训练损失下降很慢或不下降可能原因学习率太大或太小 → 尝试 0.01、0.001、0.0001数据没有归一化但我们用了 ToTensor() 已归一化到[0,1]模型太简单数据太难 → 考虑加深网络或使用预训练模型Q4验证损失一直比训练损失高很多可能过拟合增加数据增强、减小模型复杂度、增加Dropout层、增加正则化权重衰减。十、总结与下一步恭喜你你已经完整实现了一个能用的图像分类系统。回顾一下我们做了哪些事理解了CNN的基本模块卷积、池化、全连接搭建了项目结构配置、数据、模型、训练分开管理写了一个自定义数据集类读取图片和CSV标签构建了一个小型CNN模型两卷积两池化一全连接实现了训练和验证循环学会了用PyTorch训练模型保存并加载模型用于后续预测你现在可以尝试的改进方向改进点如何做预期效果增加数据增强在transform中加入T.RandomHorizontalFlip(),T.ColorJitter()提升泛化能力减少过拟合加深网络再加一组Conv2d Pool可能提高准确率但也可能过拟合使用预训练模型用torchvision.models.resnet18(pretrainedTrue)替换自己的模型大幅提升准确率尤其小数据集添加准确率指标在验证函数中计算accuracy (pred labels).float().mean()更直观地评估模型绘制混淆矩阵用sklearn.metrics.confusion_matrix看清哪几类容易混淆最后送你一句话看懂这篇文章只需要耐心但真正学会需要你亲手敲一遍代码。不要复制粘贴逐行打出来你会发现自己不知不觉就入门了深度学习的实战。如果你在运行中遇到任何问题欢迎在评论区留言。祝你在AI之路上一帆风顺