import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms, models import matplotlib.pyplot as plt库作用os操作系统接口用于创建文件夹、拼接路径、检查文件是否存在torchPyTorch 核心库提供张量运算、GPU 加速和深度学习基础功能torch.nn神经网络模块包含各种层全连接层、卷积层和损失函数torch.optim优化器模块用于更新模型权重降低损失DataLoader数据加载器将数据集分成批次批量喂给模型训练datasets包含各种内置数据集和通用数据集类如 ImageFoldertransforms图像预处理工具用于调整大小、转换张量、数据增强models包含各种预训练好的经典模型如 ResNet、VGGmatplotlib.pyplot绘图库用于绘制训练过程的损失和准确率曲线train.py全过程1. 导入所需库python运行import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms, models import matplotlib.pyplot as plt表格库作用os操作系统接口用于创建文件夹、拼接路径、检查文件是否存在torchPyTorch 核心库提供张量运算、GPU 加速和深度学习基础功能torch.nn神经网络模块包含各种层全连接层、卷积层和损失函数torch.optim优化器模块用于更新模型权重降低损失DataLoader数据加载器将数据集分成批次批量喂给模型训练datasets包含各种内置数据集和通用数据集类如 ImageFoldertransforms图像预处理工具用于调整大小、转换张量、数据增强models包含各种预训练好的经典模型如 ResNet、VGGmatplotlib.pyplot绘图库用于绘制训练过程的损失和准确率曲线2. 全局配置超参数python运行DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) IMG_SIZE 224 BATCH_SIZE 8 EPOCHS 10 LR 1e-4 NUM_CLASSES 7 FLOWER_CLASSES [ 康乃馨, 杜鹃花, 桂花, 洛神花, 荷花, 菊花, 风信子 ]表格参数作用DEVICE自动检测是否有 NVIDIA GPU有则使用 GPU 加速训练速度快 10-100 倍没有则使用 CPUIMG_SIZE输入模型的图片统一尺寸ResNet18 模型默认使用 224×224 像素BATCH_SIZE批次大小每次训练同时处理的图片数量。显存越大可以设得越大显存不足时调小如 4EPOCHS训练轮数整个训练集被模型完整遍历的次数。次数太少模型学不好太多容易过拟合LR学习率控制模型权重更新的步长。太大容易震荡不收敛太小收敛太慢NUM_CLASSES分类任务的总类别数这里是 7 种花卉必须和实际类别数完全一致FLOWER_CLASSES花卉类别名称列表必须和dataset/train下的文件夹名称完全一致用于将数字标签映射为中文名称3. 数据预处理与增强python运行train_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(20), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])transforms.Compose()将多个图像变换操作组合成一个流水线按顺序执行训练集特有增强操作RandomHorizontalFlip()随机水平翻转图片增加数据多样性让模型对左右翻转的花卉同样敏感RandomRotation(20)随机旋转图片 ±20 度让模型对不同角度的花卉更鲁棒通用操作Resize()将所有图片统一调整为 224×224 像素模型只能接受固定尺寸的输入ToTensor()将 PIL 图片0-255 整数转换为 PyTorch 张量0-1 浮点数这是模型能处理的格式Normalize()使用 ImageNet 数据集的均值和标准差进行归一化。因为我们用的 ResNet18 是在 ImageNet 上预训练的必须使用相同的归一化参数才能获得最佳效果为什么测试集没有数据增强因为测试集用于评估模型的真实性能需要使用原始图片不能做任何随机变换。4. 加载数据集python运行def load_data(): train_dataset datasets.ImageFolder(./dataset/train, transformtrain_transform) test_dataset datasets.ImageFolder(./dataset/test, transformtest_transform) print(\n✅ 类别映射关系) print(train_dataset.class_to_idx) print(\n你的FLOWER_CLASSES顺序) print(FLOWER_CLASSES) if train_dataset.classes ! FLOWER_CLASSES: print(\n⚠️ 警告数据集类别顺序与FLOWER_CLASSES不一致) print(请调整FLOWER_CLASSES顺序使其与上面的class_to_idx一致) train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, shuffleTrue, num_workers0) test_loader DataLoader(test_dataset, batch_sizeBATCH_SIZE, shuffleFalse, num_workers0) return train_loader, test_loaderdatasets.ImageFolder()PyTorch 最常用的数据集类自动按照文件夹名称分配标签例如dataset/train/康乃馨/下的所有图片都会被标记为数字0dataset/train/杜鹃花/下的所有图片都会被标记为数字1标签顺序按照文件夹名称的拼音首字母排序train_dataset.class_to_idx打印类别到数字标签的映射关系用于验证是否正确DataLoader将数据集包装成可迭代的批次加载器shuffleTrue训练集打乱顺序避免模型学习到数据的顺序特征shuffleFalse测试集不打乱方便评估结果num_workers0Windows 系统下必须设为 0否则会出现多进程错误5. 构建模型ResNet18 迁移学习python运行def build_model(): model models.resnet18(weightsNone) for param in model.parameters(): param.requires_grad False model.fc nn.Sequential( nn.Linear(model.fc.in_features, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, NUM_CLASSES) ) return model.to(DEVICE)迁移学习原理ResNet18 是一个在 ImageNet 数据集1400 万张图片1000 个类别上训练好的深度卷积神经网络已经学习到了通用的图像特征边缘、纹理、形状等。我们不需要从头训练整个网络只需要替换最后的分类层让它学会区分我们的 7 种花卉。models.resnet18(weightsNone)加载 ResNet18 模型结构不加载预训练权重因为你之前网络有问题无法下载。如果网络正常设为weightsIMAGENET1K_V1可以加载预训练权重大幅提高准确率和训练速度。param.requires_grad False冻结 ResNet18 主干网络的所有参数训练时不更新这些层的权重只训练我们添加的分类层。这样可以大大减少训练参数数量加快训练速度防止过拟合。model.fc替换 ResNet18 原有的全连接层原先是 1000 类分类nn.Linear(model.fc.in_features, 256)将 ResNet18 输出的 512 维特征降维到 256 维nn.ReLU()激活函数引入非线性让模型能够学习复杂的特征nn.Dropout(0.5)随机失活 50% 的神经元防止模型过拟合在小数据集上非常重要nn.Linear(256, NUM_CLASSES)最终输出 7 个类别的预测分数model.to(DEVICE)将模型移动到 GPU 或 CPU 上6. 单轮训练函数python运行def train(model, train_loader, criterion, optimizer): model.train() total_loss 0 correct 0 for images, labels in train_loader: images, labels images.to(DEVICE), labels.to(DEVICE) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() avg_loss total_loss / len(train_loader) acc correct / len(train_loader.dataset) return avg_loss, acc这是深度学习训练的标准四步流程前向传播将图片输入模型得到预测结果计算损失计算预测结果和真实标签之间的差距反向传播根据损失计算每个参数的梯度更新参数使用优化器更新模型参数降低损失model.train()将模型设置为训练模式启用 Dropout 和 BatchNorm 层images, labels.to(DEVICE)将数据移动到 GPU 或 CPU 上outputs model(images)前向传播得到模型对每个类别的预测分数loss criterion(outputs, labels)计算损失损失越小说明模型预测越准确optimizer.zero_grad()清空上一轮的梯度避免梯度累积loss.backward()反向传播自动计算每个参数的梯度optimizer.step()根据梯度更新模型参数loss.item()将张量转换为 Python 标量避免内存泄漏outputs.argmax(1)取每个样本预测分数最高的类别索引avg_loss本轮训练的平均损失acc本轮训练的准确率正确预测的样本数 / 总样本数7. 单轮测试函数python运行def test(model, test_loader, criterion): model.eval() total_loss 0 correct 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(DEVICE), labels.to(DEVICE) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() avg_loss total_loss / len(test_loader) acc correct / len(test_loader.dataset) return avg_loss, accmodel.eval()将模型设置为评估模式禁用 Dropout 和 BatchNorm 层torch.no_grad()非常重要测试时不需要计算梯度可以大幅节省显存和提高速度测试流程和训练类似但没有反向传播和参数更新步骤只进行前向传播和指标计算8. 绘制训练曲线python运行def plot_curve(train_loss, train_acc, test_loss, test_acc): plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(train_loss, label训练损失) plt.plot(test_loss, label测试损失) plt.title(损失曲线) plt.legend() plt.subplot(1, 2, 2) plt.plot(train_acc, label训练准确率) plt.plot(test_acc, label测试准确率) plt.title(准确率曲线) plt.legend() plt.savefig(./models/train_curve.png) plt.show()绘制训练过程中损失和准确率的变化趋势帮助判断模型的训练状态理想情况训练损失和测试损失都持续下降最终趋于稳定训练准确率和测试准确率都持续上升最终趋于稳定过拟合训练准确率很高但测试准确率很低说明模型记住了训练集的图片但泛化能力差欠拟合训练准确率和测试准确率都很低说明模型没有学到足够的特征9. 主训练循环程序入口python运行if __name__ __main__: print( * 60) print(PyTorch 花卉识别训练程序) print( * 60) print(f使用设备: {DEVICE}) print(f类别数量: {NUM_CLASSES}) os.makedirs(./models, exist_okTrue) train_loader, test_loader load_data() print(f\n✅ 数据集加载完成) print(f训练集批次数量: {len(train_loader)}) print(f测试集批次数量: {len(test_loader)}) model build_model() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.fc.parameters(), lrLR) train_loss_list, train_acc_list [], [] test_loss_list, test_acc_list [], [] best_acc 0.0 print(\n 开始训练...) for epoch in range(EPOCHS): train_loss, train_acc train(model, train_loader, criterion, optimizer) test_loss, test_acc test(model, test_loader, criterion) train_loss_list.append(train_loss) train_acc_list.append(train_acc) test_loss_list.append(test_loss) test_acc_list.append(test_acc) print( fEpoch [{epoch 1}/{EPOCHS}] | 训练损失:{train_loss:.3f} | 训练准确率:{train_acc:.3f} | 测试准确率:{test_acc:.3f}) if test_acc best_acc: best_acc test_acc torch.save(model.state_dict(), ./models/flower_model.pth) print(✅ 保存最优模型) plot_curve(train_loss_list, train_acc_list, test_loss_list, test_acc_list) print(f\n 训练完成最佳测试准确率: {best_acc:.3f})os.makedirs(./models, exist_okTrue)创建模型保存文件夹如果已经存在则不报错criterion nn.CrossEntropyLoss()多分类任务的标准损失函数内部已经包含了 Softmax 激活optimizer optim.Adam(model.fc.parameters(), lrLR)Adam 优化器自适应学习率训练效果好且稳定。这里只优化分类层的参数因为主干网络已经被冻结了。for epoch in range(EPOCHS)循环训练 EPOCHS 轮torch.save(model.state_dict(), ./models/flower_model.pth)只保存模型的权重参数推荐方式文件小且灵活为什么保存最优模型而不是最后一轮的模型因为训练后期模型可能会过拟合测试准确率会下降。保存测试准确率最高的模型可以获得最好的泛化能力。数据准备将花卉图片按类别放入dataset/train和dataset/test对应的文件夹中训练模型运行train.py模型通过学习训练集图片不断调整权重找到最优参数保存模型训练过程中自动保存测试准确率最高的模型权重到models/flower_model.pth预测识别运行predict.py加载训练好的权重对predict文件夹中的新图片进行识别训练好的模型可以反复使用以后只需要运行predict.py就能识别新的花卉图片了。predict.py 逐行详细解释这是推理阶段的代码作用是加载训练好的模型权重对新的、从未见过的花卉图片进行识别。它和train.py是强关联的很多地方必须和训练代码完全一致否则预测结果会完全错误。一、导入所需库python运行import os import torch import torch.nn as nn from torchvision import transforms, models from PIL import Image表格库作用为什么 predict 需要它os文件系统操作检查模型文件是否存在、扫描 predict 文件夹里的图片torchPyTorch 核心库加载模型权重、张量运算、GPU 加速torch.nn神经网络模块搭建和训练时完全一致的模型结构transforms图像预处理对预测图片做和训练时完全相同的预处理models预训练模型库加载 ResNet18 模型骨架PIL.Image图片加载库打开本地磁盘上的单张图片train 用 ImageFolder 自动加载predict 需要手动加载二、全局配置参数python运行# ---------------------- 配置参数 ---------------------- DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) # 注意如果训练时保存的是 best_flower_model.pth这里要改文件名 MODEL_PATH ./models/flower_model.pth PREDICT_PATH ./predict IMG_SIZE 224 # 必须和 train.py 中的类别顺序完全一致 CLASS_NAMES [ 康乃馨, 杜鹃花, 桂花, 洛神花, 荷花, 菊花, 风信子 ]逐行解释DEVICE torch.device(cuda if torch.cuda.is_available() else cpu)和 train.py 完全一样自动选择 GPU 或 CPU 运行推理时 GPU 速度也会比 CPU 快很多但 CPU 也能正常运行MODEL_PATH ./models/flower_model.pth训练好的权重文件路径必须和 train.py 中 torch.save () 保存的文件名完全一致如果训练时保存的是best_flower_model.pth这里也要改成对应的名字PREDICT_PATH ./predict待识别图片存放的文件夹路径程序会自动扫描这个文件夹里的所有图片进行批量识别IMG_SIZE 224输入图片的尺寸必须和 train.py 中的 IMG_SIZE 完全一致ResNet18 默认使用 224×224改了会导致模型输入维度不匹配CLASS_NAMES [...]类别名称列表将模型输出的数字标签映射为中文名称最关键必须和 train.py 中的 FLOWER_CLASSES 顺序、文字完全一致顺序错了会导致 指鹿为马模型预测出标签 0本应是康乃馨却会被错误映射成其他花三、图像预处理python运行# 图像预处理必须和训练时的验证集预处理完全一致 transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])逐行解释transforms.Compose([...])将多个预处理操作组合成一个流水线按顺序执行transforms.Resize((IMG_SIZE, IMG_SIZE))将任意尺寸的输入图片统一调整为 224×224 像素模型只能接受固定尺寸的输入transforms.ToTensor()将 PIL 图片0-255 整数转换为 PyTorch 张量0-1 浮点数同时将通道顺序从(H, W, C)转换为(C, H, W)这是 PyTorch 要求的格式transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])使用 ImageNet 数据集的均值和标准差对图片进行归一化必须和 train.py 中的 test_transform 完全一致模型是在这个归一化后的数值分布上训练的如果预处理不一样输入分布变了模型就认不出来了重要预测时不能用训练时的数据增强随机翻转、旋转因为我们需要对原始图片进行稳定的识别。四、模型加载函数python运行# ---------------------- 模型加载 ---------------------- def load_model(): # 1. 搭建空模型骨架不下载官方预训练权重 model models.resnet18(weightsNone) # 2. 替换全连接层必须和 train.py 结构完全一致 model.fc nn.Sequential( nn.Linear(model.fc.in_features, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, 7) # 7个类别 ) # 3. 加载本地训练好的权重 if not os.path.exists(MODEL_PATH): raise FileNotFoundError(f权重文件不存在: {MODEL_PATH}) model.load_state_dict(torch.load(MODEL_PATH, map_locationDEVICE)) model.to(DEVICE) model.eval() # 切换到推理模式关闭Dropout和BatchNorm的训练行为 return model逐行解释model models.resnet18(weightsNone)搭建一个空的 ResNet18 模型骨架不加载任何预训练权重必须和 train.py 中的模型初始化完全一致model.fc nn.Sequential(...)替换 ResNet18 原有的全连接层必须和 train.py 中替换的分类层结构完全一致包括层数、神经元数量、激活函数、Dropout 比例一个都不能错因为torch.save()只保存了模型的权重参数没有保存模型结构所以必须先建一个一模一样的空架子再把权重装进去if not os.path.exists(MODEL_PATH): raise FileNotFoundError(...)检查权重文件是否存在如果不存在抛出明确的错误提示告诉用户需要先运行 train.py 训练模型model.load_state_dict(torch.load(MODEL_PATH, map_locationDEVICE))加载本地训练好的权重文件map_locationDEVICE确保权重被加载到正确的设备GPU/CPU上如果训练时用 GPU 保存的权重推理时用 CPU 加载这个参数会自动转换model.to(DEVICE)将整个模型移动到 GPU 或 CPU 上model.eval()推理时必须调用将模型切换到评估模式作用关闭 Dropout 层不再随机失活神经元关闭 BatchNorm 层的训练模式使用训练时计算好的均值和方差而不是当前批次的统计量如果不调用model.eval()预测结果会非常不稳定每次运行同一张图片可能得到不同的结果五、单张图片预测函数python运行# ---------------------- 单张图片预测 ---------------------- def predict_image(model, img_path): try: # 打开图片并转换为RGB处理透明PNG等格式 img Image.open(img_path).convert(RGB) # 预处理 增加batch维度 img_tensor transform(img).unsqueeze(0).to(DEVICE) # 推理关闭梯度计算提速省显存 with torch.no_grad(): output model(img_tensor) prob torch.softmax(output, dim1) # 转换为概率 pred_idx torch.argmax(prob).item() confidence prob[0][pred_idx].item() * 100 # 置信度百分比 return CLASS_NAMES[pred_idx], confidence except Exception as e: return f识别失败: {str(e)}, 0.0逐行解释try: ... except Exception as e:异常处理如果某张图片损坏、无法打开程序不会崩溃只会返回识别失败的提示img Image.open(img_path).convert(RGB)打开本地磁盘上的图片文件.convert(RGB)非常重要将图片转换为 RGB 三通道格式解决的问题PNG 图片是 RGBA 四通道带透明通道模型只接受三通道输入灰度图是单通道也需要转换为三通道如果不转换会出现通道数不匹配的错误img_tensor transform(img).unsqueeze(0).to(DEVICE)transform(img)对图片进行预处理得到形状为(3, 224, 224)的张量.unsqueeze(0)必须加在最前面增加一个 batch 维度形状变为(1, 3, 224, 224)因为模型的输入要求是(batch_size, channels, height, width)即使只有一张图片也需要有 batch 维度.to(DEVICE)将图片张量移动到 GPU 或 CPU 上with torch.no_grad():推理时必须加关闭梯度计算作用大幅节省显存不需要存储计算梯度所需的中间结果提高推理速度不需要计算梯度推理时我们只需要前向传播的结果不需要反向传播所以完全不需要梯度output model(img_tensor)前向传播将预处理后的图片输入模型得到输出输出形状为(1, 7)7 个数值分别对应 7 类花卉的原始得分不是概率prob torch.softmax(output, dim1)将模型输出的原始得分转换为概率Softmax 函数会将所有数值映射到 0-1 之间并且所有类别的概率之和为 1dim1在类别维度上进行 Softmax 计算pred_idx torch.argmax(prob).item()取概率最大的那个类别的索引.item()将 PyTorch 张量转换为 Python 整数confidence prob[0][pred_idx].item() * 100计算置信度转换为百分比形式例如0.95 → 95%表示模型有 95% 的把握认为这张图片是该类别return CLASS_NAMES[pred_idx], confidence根据索引从 CLASS_NAMES 列表中取出对应的中文名称返回识别结果和置信度except Exception as e: return f识别失败: {str(e)}, 0.0如果出现任何异常图片损坏、路径错误等返回识别失败的提示和 0 置信度六、批量预测文件夹函数python运行# ---------------------- 批量预测文件夹 ---------------------- def batch_predict(model): # 自动创建预测文件夹 if not os.path.exists(PREDICT_PATH): os.makedirs(PREDICT_PATH) print(f✅ 已创建预测文件夹: {PREDICT_PATH}) print(请将需要识别的花卉图片放入该文件夹后重新运行) return # 筛选支持的图片格式 img_extensions (.jpg, .jpeg, .png, .bmp, .webp) img_files [f for f in os.listdir(PREDICT_PATH) if f.lower().endswith(img_extensions)] if not img_files: print(f❌ 预测文件夹 {PREDICT_PATH} 中没有找到图片) return print(f\n 找到 {len(img_files)} 张图片开始识别...) print(- * 60) for filename in img_files: full_path os.path.join(PREDICT_PATH, filename) result, confidence predict_image(model, full_path) print(f图片: {filename}) print(f识别结果: {result}) print(f置信度: {confidence:.2f}%) print(- * 60)逐行解释if not os.path.exists(PREDICT_PATH): os.makedirs(PREDICT_PATH)如果 predict 文件夹不存在自动创建它并提示用户将需要识别的图片放入该文件夹img_extensions (.jpg, .jpeg, .png, .bmp, .webp)定义支持的图片格式可以根据需要添加更多格式img_files [f for f in os.listdir(PREDICT_PATH) if f.lower().endswith(img_extensions)]扫描 predict 文件夹中的所有文件筛选出后缀在支持列表中的图片文件.lower()将文件名转为小写避免因为后缀大小写问题如.JPG被漏掉if not img_files: print(没有找到图片)如果文件夹中没有图片给出提示并返回for filename in img_files:遍历所有找到的图片文件full_path os.path.join(PREDICT_PATH, filename)拼接图片的完整路径使用os.path.join()保证跨平台兼容性result, confidence predict_image(model, full_path)调用单张图片预测函数对当前图片进行识别print(f图片: {filename}) ...打印识别结果和置信度confidence:.2f%将置信度保留两位小数七、程序入口主函数python运行if __name__ __main__: print( * 60) print(PyTorch 花卉识别程序) print( * 60) try: model load_model() print(✅ 模型加载成功) except Exception as e: print(f❌ 模型加载失败: {e}) print(请先运行train.py训练模型并确保models文件夹中有flower_model.pth文件) exit() batch_predict(model) print(\n 所有图片识别完成)逐行解释if __name__ __main__:Python 标准程序入口只有当直接运行这个文件时才会执行下面的代码try: model load_model() except Exception as e:尝试加载模型如果加载失败权重文件不存在、模型结构不匹配等打印详细错误信息并退出程序batch_predict(model)调用批量预测函数对 predict 文件夹中的所有图片进行识别print(\n 所有图片识别完成)所有图片识别完成后给出提示八、关键注意事项总结所有和 train.py 相关的配置必须完全一致CLASS_NAMES 的顺序和文字IMG_SIZE预处理的参数特别是 Normalize 的均值和标准差模型结构替换的全连接层推理时必须做的三件事调用model.eval()切换到评估模式使用with torch.no_grad()关闭梯度计算预处理和训练时的 test_transform 完全一致常见错误模型结构和训练时不一样 → 加载权重时报错CLASS_NAMES 顺序不对 → 识别结果错位预处理不一样 → 识别准确率极低没有加unsqueeze(0)→ 维度不匹配错误没有调用model.eval()→ 结果不稳定
深度学习花卉识别笔记
import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms, models import matplotlib.pyplot as plt库作用os操作系统接口用于创建文件夹、拼接路径、检查文件是否存在torchPyTorch 核心库提供张量运算、GPU 加速和深度学习基础功能torch.nn神经网络模块包含各种层全连接层、卷积层和损失函数torch.optim优化器模块用于更新模型权重降低损失DataLoader数据加载器将数据集分成批次批量喂给模型训练datasets包含各种内置数据集和通用数据集类如 ImageFoldertransforms图像预处理工具用于调整大小、转换张量、数据增强models包含各种预训练好的经典模型如 ResNet、VGGmatplotlib.pyplot绘图库用于绘制训练过程的损失和准确率曲线train.py全过程1. 导入所需库python运行import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms, models import matplotlib.pyplot as plt表格库作用os操作系统接口用于创建文件夹、拼接路径、检查文件是否存在torchPyTorch 核心库提供张量运算、GPU 加速和深度学习基础功能torch.nn神经网络模块包含各种层全连接层、卷积层和损失函数torch.optim优化器模块用于更新模型权重降低损失DataLoader数据加载器将数据集分成批次批量喂给模型训练datasets包含各种内置数据集和通用数据集类如 ImageFoldertransforms图像预处理工具用于调整大小、转换张量、数据增强models包含各种预训练好的经典模型如 ResNet、VGGmatplotlib.pyplot绘图库用于绘制训练过程的损失和准确率曲线2. 全局配置超参数python运行DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) IMG_SIZE 224 BATCH_SIZE 8 EPOCHS 10 LR 1e-4 NUM_CLASSES 7 FLOWER_CLASSES [ 康乃馨, 杜鹃花, 桂花, 洛神花, 荷花, 菊花, 风信子 ]表格参数作用DEVICE自动检测是否有 NVIDIA GPU有则使用 GPU 加速训练速度快 10-100 倍没有则使用 CPUIMG_SIZE输入模型的图片统一尺寸ResNet18 模型默认使用 224×224 像素BATCH_SIZE批次大小每次训练同时处理的图片数量。显存越大可以设得越大显存不足时调小如 4EPOCHS训练轮数整个训练集被模型完整遍历的次数。次数太少模型学不好太多容易过拟合LR学习率控制模型权重更新的步长。太大容易震荡不收敛太小收敛太慢NUM_CLASSES分类任务的总类别数这里是 7 种花卉必须和实际类别数完全一致FLOWER_CLASSES花卉类别名称列表必须和dataset/train下的文件夹名称完全一致用于将数字标签映射为中文名称3. 数据预处理与增强python运行train_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(20), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])transforms.Compose()将多个图像变换操作组合成一个流水线按顺序执行训练集特有增强操作RandomHorizontalFlip()随机水平翻转图片增加数据多样性让模型对左右翻转的花卉同样敏感RandomRotation(20)随机旋转图片 ±20 度让模型对不同角度的花卉更鲁棒通用操作Resize()将所有图片统一调整为 224×224 像素模型只能接受固定尺寸的输入ToTensor()将 PIL 图片0-255 整数转换为 PyTorch 张量0-1 浮点数这是模型能处理的格式Normalize()使用 ImageNet 数据集的均值和标准差进行归一化。因为我们用的 ResNet18 是在 ImageNet 上预训练的必须使用相同的归一化参数才能获得最佳效果为什么测试集没有数据增强因为测试集用于评估模型的真实性能需要使用原始图片不能做任何随机变换。4. 加载数据集python运行def load_data(): train_dataset datasets.ImageFolder(./dataset/train, transformtrain_transform) test_dataset datasets.ImageFolder(./dataset/test, transformtest_transform) print(\n✅ 类别映射关系) print(train_dataset.class_to_idx) print(\n你的FLOWER_CLASSES顺序) print(FLOWER_CLASSES) if train_dataset.classes ! FLOWER_CLASSES: print(\n⚠️ 警告数据集类别顺序与FLOWER_CLASSES不一致) print(请调整FLOWER_CLASSES顺序使其与上面的class_to_idx一致) train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, shuffleTrue, num_workers0) test_loader DataLoader(test_dataset, batch_sizeBATCH_SIZE, shuffleFalse, num_workers0) return train_loader, test_loaderdatasets.ImageFolder()PyTorch 最常用的数据集类自动按照文件夹名称分配标签例如dataset/train/康乃馨/下的所有图片都会被标记为数字0dataset/train/杜鹃花/下的所有图片都会被标记为数字1标签顺序按照文件夹名称的拼音首字母排序train_dataset.class_to_idx打印类别到数字标签的映射关系用于验证是否正确DataLoader将数据集包装成可迭代的批次加载器shuffleTrue训练集打乱顺序避免模型学习到数据的顺序特征shuffleFalse测试集不打乱方便评估结果num_workers0Windows 系统下必须设为 0否则会出现多进程错误5. 构建模型ResNet18 迁移学习python运行def build_model(): model models.resnet18(weightsNone) for param in model.parameters(): param.requires_grad False model.fc nn.Sequential( nn.Linear(model.fc.in_features, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, NUM_CLASSES) ) return model.to(DEVICE)迁移学习原理ResNet18 是一个在 ImageNet 数据集1400 万张图片1000 个类别上训练好的深度卷积神经网络已经学习到了通用的图像特征边缘、纹理、形状等。我们不需要从头训练整个网络只需要替换最后的分类层让它学会区分我们的 7 种花卉。models.resnet18(weightsNone)加载 ResNet18 模型结构不加载预训练权重因为你之前网络有问题无法下载。如果网络正常设为weightsIMAGENET1K_V1可以加载预训练权重大幅提高准确率和训练速度。param.requires_grad False冻结 ResNet18 主干网络的所有参数训练时不更新这些层的权重只训练我们添加的分类层。这样可以大大减少训练参数数量加快训练速度防止过拟合。model.fc替换 ResNet18 原有的全连接层原先是 1000 类分类nn.Linear(model.fc.in_features, 256)将 ResNet18 输出的 512 维特征降维到 256 维nn.ReLU()激活函数引入非线性让模型能够学习复杂的特征nn.Dropout(0.5)随机失活 50% 的神经元防止模型过拟合在小数据集上非常重要nn.Linear(256, NUM_CLASSES)最终输出 7 个类别的预测分数model.to(DEVICE)将模型移动到 GPU 或 CPU 上6. 单轮训练函数python运行def train(model, train_loader, criterion, optimizer): model.train() total_loss 0 correct 0 for images, labels in train_loader: images, labels images.to(DEVICE), labels.to(DEVICE) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() avg_loss total_loss / len(train_loader) acc correct / len(train_loader.dataset) return avg_loss, acc这是深度学习训练的标准四步流程前向传播将图片输入模型得到预测结果计算损失计算预测结果和真实标签之间的差距反向传播根据损失计算每个参数的梯度更新参数使用优化器更新模型参数降低损失model.train()将模型设置为训练模式启用 Dropout 和 BatchNorm 层images, labels.to(DEVICE)将数据移动到 GPU 或 CPU 上outputs model(images)前向传播得到模型对每个类别的预测分数loss criterion(outputs, labels)计算损失损失越小说明模型预测越准确optimizer.zero_grad()清空上一轮的梯度避免梯度累积loss.backward()反向传播自动计算每个参数的梯度optimizer.step()根据梯度更新模型参数loss.item()将张量转换为 Python 标量避免内存泄漏outputs.argmax(1)取每个样本预测分数最高的类别索引avg_loss本轮训练的平均损失acc本轮训练的准确率正确预测的样本数 / 总样本数7. 单轮测试函数python运行def test(model, test_loader, criterion): model.eval() total_loss 0 correct 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(DEVICE), labels.to(DEVICE) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() avg_loss total_loss / len(test_loader) acc correct / len(test_loader.dataset) return avg_loss, accmodel.eval()将模型设置为评估模式禁用 Dropout 和 BatchNorm 层torch.no_grad()非常重要测试时不需要计算梯度可以大幅节省显存和提高速度测试流程和训练类似但没有反向传播和参数更新步骤只进行前向传播和指标计算8. 绘制训练曲线python运行def plot_curve(train_loss, train_acc, test_loss, test_acc): plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(train_loss, label训练损失) plt.plot(test_loss, label测试损失) plt.title(损失曲线) plt.legend() plt.subplot(1, 2, 2) plt.plot(train_acc, label训练准确率) plt.plot(test_acc, label测试准确率) plt.title(准确率曲线) plt.legend() plt.savefig(./models/train_curve.png) plt.show()绘制训练过程中损失和准确率的变化趋势帮助判断模型的训练状态理想情况训练损失和测试损失都持续下降最终趋于稳定训练准确率和测试准确率都持续上升最终趋于稳定过拟合训练准确率很高但测试准确率很低说明模型记住了训练集的图片但泛化能力差欠拟合训练准确率和测试准确率都很低说明模型没有学到足够的特征9. 主训练循环程序入口python运行if __name__ __main__: print( * 60) print(PyTorch 花卉识别训练程序) print( * 60) print(f使用设备: {DEVICE}) print(f类别数量: {NUM_CLASSES}) os.makedirs(./models, exist_okTrue) train_loader, test_loader load_data() print(f\n✅ 数据集加载完成) print(f训练集批次数量: {len(train_loader)}) print(f测试集批次数量: {len(test_loader)}) model build_model() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.fc.parameters(), lrLR) train_loss_list, train_acc_list [], [] test_loss_list, test_acc_list [], [] best_acc 0.0 print(\n 开始训练...) for epoch in range(EPOCHS): train_loss, train_acc train(model, train_loader, criterion, optimizer) test_loss, test_acc test(model, test_loader, criterion) train_loss_list.append(train_loss) train_acc_list.append(train_acc) test_loss_list.append(test_loss) test_acc_list.append(test_acc) print( fEpoch [{epoch 1}/{EPOCHS}] | 训练损失:{train_loss:.3f} | 训练准确率:{train_acc:.3f} | 测试准确率:{test_acc:.3f}) if test_acc best_acc: best_acc test_acc torch.save(model.state_dict(), ./models/flower_model.pth) print(✅ 保存最优模型) plot_curve(train_loss_list, train_acc_list, test_loss_list, test_acc_list) print(f\n 训练完成最佳测试准确率: {best_acc:.3f})os.makedirs(./models, exist_okTrue)创建模型保存文件夹如果已经存在则不报错criterion nn.CrossEntropyLoss()多分类任务的标准损失函数内部已经包含了 Softmax 激活optimizer optim.Adam(model.fc.parameters(), lrLR)Adam 优化器自适应学习率训练效果好且稳定。这里只优化分类层的参数因为主干网络已经被冻结了。for epoch in range(EPOCHS)循环训练 EPOCHS 轮torch.save(model.state_dict(), ./models/flower_model.pth)只保存模型的权重参数推荐方式文件小且灵活为什么保存最优模型而不是最后一轮的模型因为训练后期模型可能会过拟合测试准确率会下降。保存测试准确率最高的模型可以获得最好的泛化能力。数据准备将花卉图片按类别放入dataset/train和dataset/test对应的文件夹中训练模型运行train.py模型通过学习训练集图片不断调整权重找到最优参数保存模型训练过程中自动保存测试准确率最高的模型权重到models/flower_model.pth预测识别运行predict.py加载训练好的权重对predict文件夹中的新图片进行识别训练好的模型可以反复使用以后只需要运行predict.py就能识别新的花卉图片了。predict.py 逐行详细解释这是推理阶段的代码作用是加载训练好的模型权重对新的、从未见过的花卉图片进行识别。它和train.py是强关联的很多地方必须和训练代码完全一致否则预测结果会完全错误。一、导入所需库python运行import os import torch import torch.nn as nn from torchvision import transforms, models from PIL import Image表格库作用为什么 predict 需要它os文件系统操作检查模型文件是否存在、扫描 predict 文件夹里的图片torchPyTorch 核心库加载模型权重、张量运算、GPU 加速torch.nn神经网络模块搭建和训练时完全一致的模型结构transforms图像预处理对预测图片做和训练时完全相同的预处理models预训练模型库加载 ResNet18 模型骨架PIL.Image图片加载库打开本地磁盘上的单张图片train 用 ImageFolder 自动加载predict 需要手动加载二、全局配置参数python运行# ---------------------- 配置参数 ---------------------- DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) # 注意如果训练时保存的是 best_flower_model.pth这里要改文件名 MODEL_PATH ./models/flower_model.pth PREDICT_PATH ./predict IMG_SIZE 224 # 必须和 train.py 中的类别顺序完全一致 CLASS_NAMES [ 康乃馨, 杜鹃花, 桂花, 洛神花, 荷花, 菊花, 风信子 ]逐行解释DEVICE torch.device(cuda if torch.cuda.is_available() else cpu)和 train.py 完全一样自动选择 GPU 或 CPU 运行推理时 GPU 速度也会比 CPU 快很多但 CPU 也能正常运行MODEL_PATH ./models/flower_model.pth训练好的权重文件路径必须和 train.py 中 torch.save () 保存的文件名完全一致如果训练时保存的是best_flower_model.pth这里也要改成对应的名字PREDICT_PATH ./predict待识别图片存放的文件夹路径程序会自动扫描这个文件夹里的所有图片进行批量识别IMG_SIZE 224输入图片的尺寸必须和 train.py 中的 IMG_SIZE 完全一致ResNet18 默认使用 224×224改了会导致模型输入维度不匹配CLASS_NAMES [...]类别名称列表将模型输出的数字标签映射为中文名称最关键必须和 train.py 中的 FLOWER_CLASSES 顺序、文字完全一致顺序错了会导致 指鹿为马模型预测出标签 0本应是康乃馨却会被错误映射成其他花三、图像预处理python运行# 图像预处理必须和训练时的验证集预处理完全一致 transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])逐行解释transforms.Compose([...])将多个预处理操作组合成一个流水线按顺序执行transforms.Resize((IMG_SIZE, IMG_SIZE))将任意尺寸的输入图片统一调整为 224×224 像素模型只能接受固定尺寸的输入transforms.ToTensor()将 PIL 图片0-255 整数转换为 PyTorch 张量0-1 浮点数同时将通道顺序从(H, W, C)转换为(C, H, W)这是 PyTorch 要求的格式transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])使用 ImageNet 数据集的均值和标准差对图片进行归一化必须和 train.py 中的 test_transform 完全一致模型是在这个归一化后的数值分布上训练的如果预处理不一样输入分布变了模型就认不出来了重要预测时不能用训练时的数据增强随机翻转、旋转因为我们需要对原始图片进行稳定的识别。四、模型加载函数python运行# ---------------------- 模型加载 ---------------------- def load_model(): # 1. 搭建空模型骨架不下载官方预训练权重 model models.resnet18(weightsNone) # 2. 替换全连接层必须和 train.py 结构完全一致 model.fc nn.Sequential( nn.Linear(model.fc.in_features, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, 7) # 7个类别 ) # 3. 加载本地训练好的权重 if not os.path.exists(MODEL_PATH): raise FileNotFoundError(f权重文件不存在: {MODEL_PATH}) model.load_state_dict(torch.load(MODEL_PATH, map_locationDEVICE)) model.to(DEVICE) model.eval() # 切换到推理模式关闭Dropout和BatchNorm的训练行为 return model逐行解释model models.resnet18(weightsNone)搭建一个空的 ResNet18 模型骨架不加载任何预训练权重必须和 train.py 中的模型初始化完全一致model.fc nn.Sequential(...)替换 ResNet18 原有的全连接层必须和 train.py 中替换的分类层结构完全一致包括层数、神经元数量、激活函数、Dropout 比例一个都不能错因为torch.save()只保存了模型的权重参数没有保存模型结构所以必须先建一个一模一样的空架子再把权重装进去if not os.path.exists(MODEL_PATH): raise FileNotFoundError(...)检查权重文件是否存在如果不存在抛出明确的错误提示告诉用户需要先运行 train.py 训练模型model.load_state_dict(torch.load(MODEL_PATH, map_locationDEVICE))加载本地训练好的权重文件map_locationDEVICE确保权重被加载到正确的设备GPU/CPU上如果训练时用 GPU 保存的权重推理时用 CPU 加载这个参数会自动转换model.to(DEVICE)将整个模型移动到 GPU 或 CPU 上model.eval()推理时必须调用将模型切换到评估模式作用关闭 Dropout 层不再随机失活神经元关闭 BatchNorm 层的训练模式使用训练时计算好的均值和方差而不是当前批次的统计量如果不调用model.eval()预测结果会非常不稳定每次运行同一张图片可能得到不同的结果五、单张图片预测函数python运行# ---------------------- 单张图片预测 ---------------------- def predict_image(model, img_path): try: # 打开图片并转换为RGB处理透明PNG等格式 img Image.open(img_path).convert(RGB) # 预处理 增加batch维度 img_tensor transform(img).unsqueeze(0).to(DEVICE) # 推理关闭梯度计算提速省显存 with torch.no_grad(): output model(img_tensor) prob torch.softmax(output, dim1) # 转换为概率 pred_idx torch.argmax(prob).item() confidence prob[0][pred_idx].item() * 100 # 置信度百分比 return CLASS_NAMES[pred_idx], confidence except Exception as e: return f识别失败: {str(e)}, 0.0逐行解释try: ... except Exception as e:异常处理如果某张图片损坏、无法打开程序不会崩溃只会返回识别失败的提示img Image.open(img_path).convert(RGB)打开本地磁盘上的图片文件.convert(RGB)非常重要将图片转换为 RGB 三通道格式解决的问题PNG 图片是 RGBA 四通道带透明通道模型只接受三通道输入灰度图是单通道也需要转换为三通道如果不转换会出现通道数不匹配的错误img_tensor transform(img).unsqueeze(0).to(DEVICE)transform(img)对图片进行预处理得到形状为(3, 224, 224)的张量.unsqueeze(0)必须加在最前面增加一个 batch 维度形状变为(1, 3, 224, 224)因为模型的输入要求是(batch_size, channels, height, width)即使只有一张图片也需要有 batch 维度.to(DEVICE)将图片张量移动到 GPU 或 CPU 上with torch.no_grad():推理时必须加关闭梯度计算作用大幅节省显存不需要存储计算梯度所需的中间结果提高推理速度不需要计算梯度推理时我们只需要前向传播的结果不需要反向传播所以完全不需要梯度output model(img_tensor)前向传播将预处理后的图片输入模型得到输出输出形状为(1, 7)7 个数值分别对应 7 类花卉的原始得分不是概率prob torch.softmax(output, dim1)将模型输出的原始得分转换为概率Softmax 函数会将所有数值映射到 0-1 之间并且所有类别的概率之和为 1dim1在类别维度上进行 Softmax 计算pred_idx torch.argmax(prob).item()取概率最大的那个类别的索引.item()将 PyTorch 张量转换为 Python 整数confidence prob[0][pred_idx].item() * 100计算置信度转换为百分比形式例如0.95 → 95%表示模型有 95% 的把握认为这张图片是该类别return CLASS_NAMES[pred_idx], confidence根据索引从 CLASS_NAMES 列表中取出对应的中文名称返回识别结果和置信度except Exception as e: return f识别失败: {str(e)}, 0.0如果出现任何异常图片损坏、路径错误等返回识别失败的提示和 0 置信度六、批量预测文件夹函数python运行# ---------------------- 批量预测文件夹 ---------------------- def batch_predict(model): # 自动创建预测文件夹 if not os.path.exists(PREDICT_PATH): os.makedirs(PREDICT_PATH) print(f✅ 已创建预测文件夹: {PREDICT_PATH}) print(请将需要识别的花卉图片放入该文件夹后重新运行) return # 筛选支持的图片格式 img_extensions (.jpg, .jpeg, .png, .bmp, .webp) img_files [f for f in os.listdir(PREDICT_PATH) if f.lower().endswith(img_extensions)] if not img_files: print(f❌ 预测文件夹 {PREDICT_PATH} 中没有找到图片) return print(f\n 找到 {len(img_files)} 张图片开始识别...) print(- * 60) for filename in img_files: full_path os.path.join(PREDICT_PATH, filename) result, confidence predict_image(model, full_path) print(f图片: {filename}) print(f识别结果: {result}) print(f置信度: {confidence:.2f}%) print(- * 60)逐行解释if not os.path.exists(PREDICT_PATH): os.makedirs(PREDICT_PATH)如果 predict 文件夹不存在自动创建它并提示用户将需要识别的图片放入该文件夹img_extensions (.jpg, .jpeg, .png, .bmp, .webp)定义支持的图片格式可以根据需要添加更多格式img_files [f for f in os.listdir(PREDICT_PATH) if f.lower().endswith(img_extensions)]扫描 predict 文件夹中的所有文件筛选出后缀在支持列表中的图片文件.lower()将文件名转为小写避免因为后缀大小写问题如.JPG被漏掉if not img_files: print(没有找到图片)如果文件夹中没有图片给出提示并返回for filename in img_files:遍历所有找到的图片文件full_path os.path.join(PREDICT_PATH, filename)拼接图片的完整路径使用os.path.join()保证跨平台兼容性result, confidence predict_image(model, full_path)调用单张图片预测函数对当前图片进行识别print(f图片: {filename}) ...打印识别结果和置信度confidence:.2f%将置信度保留两位小数七、程序入口主函数python运行if __name__ __main__: print( * 60) print(PyTorch 花卉识别程序) print( * 60) try: model load_model() print(✅ 模型加载成功) except Exception as e: print(f❌ 模型加载失败: {e}) print(请先运行train.py训练模型并确保models文件夹中有flower_model.pth文件) exit() batch_predict(model) print(\n 所有图片识别完成)逐行解释if __name__ __main__:Python 标准程序入口只有当直接运行这个文件时才会执行下面的代码try: model load_model() except Exception as e:尝试加载模型如果加载失败权重文件不存在、模型结构不匹配等打印详细错误信息并退出程序batch_predict(model)调用批量预测函数对 predict 文件夹中的所有图片进行识别print(\n 所有图片识别完成)所有图片识别完成后给出提示八、关键注意事项总结所有和 train.py 相关的配置必须完全一致CLASS_NAMES 的顺序和文字IMG_SIZE预处理的参数特别是 Normalize 的均值和标准差模型结构替换的全连接层推理时必须做的三件事调用model.eval()切换到评估模式使用with torch.no_grad()关闭梯度计算预处理和训练时的 test_transform 完全一致常见错误模型结构和训练时不一样 → 加载权重时报错CLASS_NAMES 顺序不对 → 识别结果错位预处理不一样 → 识别准确率极低没有加unsqueeze(0)→ 维度不匹配错误没有调用model.eval()→ 结果不稳定
