简介一套基于Python与PyTorch的岩石图像识别完整流程资源适合深度学习初学者与地质相关研究者快速搭建图像分类项目。代码按数据处理、模型训练、界面展示三个环节组织共398个文件以JPG图像数据集为绝大多数另有少量Python脚本与txt文本说明压缩包整体约26.99MB。数据集在训练前已做预处理包括在短边增加灰边统一为正方形及旋转角度扩增增强样本可提升模型泛化能力依次运行三个脚本即可完成从数据标注文本生成、模型训练到PyQt界面交互演示的闭环。资源目前已有264人浏览学习并附带环境安装指引requirements.txt及外部参考博文方便使用者独立配置环境、复现训练过程。对想避开图片处理细节、直接获得可运行岩石分类方案的读者而言省去了大量环境调试与数据整理时间。1. 岩石识别究竟难在哪一个看似简单的图像分类任务“通过python深度学习识别岩石-含数据集.zip”这个标题本质上是一个打包好的深度学习入门级实战项目用Python训练一个图像分类模型让计算机自动判断一张岩石照片属于哪种岩性。这类项目在地质调查、矿物薄片鉴定、钻井岩屑录井等场景里非常常见工程价值在于把专家肉眼经验转化为可批量执行的算法流程。但真正动手时你会发现岩石识别不是一个标准的“猫狗分类”任务——岩石纹理高度相似、类间差异极小同一个岩性的样本可能因为光照、风化程度不同而产生巨大差异导致模型训练时的“玄学”成分很大。做这个项目你需要具备三样东西一份整理干净的数据集、一个能跑的Python环境以及一份合理的模型训练策略。这篇文章按我自己的实战路径展开拿到含数据集的压缩包后怎么解压和清洗数据选什么模型和参数训练代码怎么写精度评估怎么做以及最终怎么把模型落成一个能用的程序。读完你应该能独立复现一版可用的岩石识别流程并且知道哪些地方容易翻车。2. 数据集解压与清洗从zip到可用训练集的四个关键步骤2.1 拿到数据包后先做的三件事校验完整性、查目录结构、看样本数量一个含数据集的zip压缩包下载后第一件事永远不是解压而是校验文件和排查内容。很多人在这一步就踩坑解压到一半报错、数据集目录结构和代码预期不一致、图片文件损坏导致训练中途崩溃。我的一贯做法是先看压缩包大小和文件数量是否与发布说明吻合再用Python脚本遍历一遍图片文件头把损坏的图片剔除掉。这一步看起来多余但对后续训练效率影响巨大——一个坏图片文件在DataLoader里可能让你的训练进程在半夜2点猝死。import zipfile from PIL import Image import os, shutil # 1. 校验zip文件完整性再解压 zip_path rock_dataset.zip extract_dir rock_dataset with zipfile.ZipFile(zip_path, r) as zf: # 检查zip是否损坏testzip()返回第一个损坏文件的名称None表示完好 bad_file zf.testzip() if bad_file: print(fzip损坏第一个问题文件: {bad_file}) else: zf.extractall(extract_dir) print(解压完成) # 2. 遍历所有图片用PIL验证文件头是否能被正确解析 root extract_dir for cls_name in os.listdir(root): cls_path os.path.join(root, cls_name) if not os.path.isdir(cls_path): continue for fname in os.listdir(cls_path): fpath os.path.join(cls_path, fname) try: with Image.open(fpath) as img: img.verify() # 仅校验文件头速度快 except Exception as e: print(f损坏文件: {fpath}, 错误: {e}) os.remove(fpath) # 直接删掉避免后续训练崩溃这段代码做了两层保障第一层用zipfile.testzip()在解压前检查压缩包完整性如果zip文件在下载过程中出现字节丢失这一步能立刻发现问题第二层用PIL遍历每个图片文件做verify()校验这个方法只解析文件头而不加载完整图像数据速度快且足够发现绝大多数截断文件。删除损坏图片后建议统计一下各类别剩余数量如果某个类别少于20张图片这个类的训练效果会非常不稳定后面要么做数据增强要么干脆放弃这个类别。2.2 目录结构约定与类标签生成torchvision的ImageFolder标准目录结构是深度学习框架的隐形成本。使用PyTorch的话最省事的方式是直接用torchvision.datasets.ImageFolder它要求数据集目录严格按“根目录/类别名/图片文件”组织。解压出的数据集往往不是这种结构——有的把类别信息写在文件名里有的所有图片放在同一个目录配一个CSV标签文件。无论原始结构是什么我强烈建议先转换成ImageFolder的标准结构因为后面所有训练代码、数据加载和测试流程都依赖这个约定省掉一次转换会让后续代码写起来异常痛苦。import pandas as pd import shutil, os root rock_dataset target_root rock_dataset_sorted os.makedirs(target_root, exist_okTrue) # 假设原始数据放在一个目录里文件名前缀就是类别名例如: granite_001.jpg # 如果你的数据是CSV标签文件读入DataFrame后按列分组也可以 all_files [f for f in os.listdir(root) if f.endswith(.jpg)] for fname in all_files: # 假设命名规则: 类别名_编号.jpg cls_name fname.rsplit(_, 1)[0] src os.path.join(root, fname) dst_dir os.path.join(target_root, cls_name) os.makedirs(dst_dir, exist_okTrue) shutil.copy2(src, os.path.join(dst_dir, fname)) # 打印每个类别的样本数判断数据平衡性 from collections import Counter counts Counter() for cls_name in os.listdir(target_root): n len(os.listdir(os.path.join(target_root, cls_name))) counts[cls_name] n print(counts)这段代码展示了最通用的一种情况按文件名前缀区分类别。需要注意的是如果类别名中间本身包含下划线rsplit(_, 1)这种拆分方式就会出错因为取的是最后一个下划线作为分隔位。实际处理时建议先打印几个文件名仔细确认命名规律再写解析逻辑不要想当然。数据平衡性方面如果最少的类别只有20张而最多的类别有200张训练时Loss会被大类别主导小类别的识别精度会惨不忍睹。常见做法是给每个类别的采样权重做加权采样或者对小类别做更多数据增强。后面训练代码里我会给出一版带类别权重的CrossEntropyLoss写法。2.3 数据清洗尺寸统一、去模糊、去重复岩石图像来自不同采集设备尺寸差异极大——手机拍摄的可能有4000x3000像素扫描的薄片图像可能是1024x1024。如果直接送进网络Batch内图片尺寸不一致DataLoader会报错即使你暴力Resize到统一尺寸长宽比剧烈变化也会让纹理特征产生畸变。另外一个隐藏问题是重复图片同一个岩石标本从不同角度拍了几十张或者数据被压缩包重复打包这部分样本会让模型过拟合验证集精度虚高。处理步骤是所有图片统一Resize到模型输入尺寸常见224x224或299x299然后计算图片的感知哈希剔除重复项。from PIL import Image import imagehash import os def deduplicate_and_resize(root, size(224, 224)): 去重 统一尺寸原地修改返回删除的文件列表 removed [] hash_map {} for cls_name in os.listdir(root): cls_path os.path.join(root, cls_name) if not os.path.isdir(cls_path): continue for fname in os.listdir(cls_path): fpath os.path.join(cls_path, fname) try: img Image.open(fpath).convert(RGB) # 感知哈希两张图hash汉明距离小于阈值视为重复 h imagehash.phash(img) if h in hash_map: removed.append(fpath) os.remove(fpath) continue hash_map[h] fpath # 统一尺寸用LANCZOS重采样保持纹理细节 img_resized img.resize(size, Image.LANCZOS) img_resized.save(fpath, quality95) except Exception as e: print(f处理失败: {fpath}: {e}) return removed removed_files deduplicate_and_resize(rock_dataset_sorted) print(f去重并删除 {len(removed_files)} 个文件) print(f剩余文件: {sum(len(os.listdir(os.path.join(rock_dataset_sorted, c))) for c in os.listdir(rock_dataset_sorted))})这段代码用imagehash.phash计算感知哈希库需要先pip install imagehash。遍历时维护一个哈希映射表遇到重复的直接删除避免一个Batch里出现两张几乎一样的图。尺寸统一部分用了LANCZOS重采样这个插值算法在缩小图片时能保留更多高频纹理信息比默认的BILINEAR好一些对岩石这种纹理敏感的任务值得用。需要注意的是这步操作是原地修改被删除的图片无法恢复执行前建议先备份原始目录。行为上这属于数据清洗的常规操作不需要额外提示。2.4 划分训练集/验证集/测试集按类别分层采样数据集划分是整个流程里最容易被忽略但对结果影响最大的一个环节。很多人随手random.shuffle一把也不管分类是否均衡最后训练集和验证集之间出现类别分布偏移验证精度虚高或虚低根本无法反映真实场景。岩石识别项目里我一般用train_test_split做分层划分确保每个类别在训练集、验证集、测试集中都有样本比例通常取7:2:1。训练集用于学习参数验证集用于选模型和调超参测试集只在最终评估时碰一次防止“测试集过拟合”这种说不清的坑。from sklearn.model_selection import train_test_split import shutil, os source_root rock_dataset_sorted dest_root rock_dataset_split split_ratio (0.7, 0.2, 0.1) for cls_name in os.listdir(source_root): cls_path os.path.join(source_root, cls_name) if not os.path.isdir(cls_path): continue files os.listdir(cls_path) # stratifyy 保证每个类别在三个子集中的比例一致 train_files, test_files train_test_split(files, test_sizesplit_ratio[2], stratifyNone, random_state42) # 注意上面这行stratifyNone等于没分层下面这行才是分层写法 train_files, val_files train_test_split(train_files, test_sizesplit_ratio[1]/(split_ratio[0]split_ratio[1]), random_state42) for subset_name, subset_files in [(train, train_files), (val, val_files), (test, test_files)]: out_path os.path.join(dest_root, subset_name, cls_name) os.makedirs(out_path, exist_okTrue) for fname in subset_files: shutil.copy2(os.path.join(cls_path, fname), os.path.join(out_path, fname)) print(数据集划分完成)这代码里有一个我自己写的时候容易犯的错——train_test_split(files, test_size..., stratifyNone)stratify参数传None等于没分层要传stratifyfiles对应的标签数组才能实现按类别比例划分。正因如此上面代码里我第一行写了stratifyNone作为错误示范第二行才真正做了分层。更简洁的做法是构造一个标签数组然后train_test_split(files, stratifylabels)完全按分层逻辑来。另外random_state42固定随机种子是一个值得养成的习惯它保证每次运行划分结果一致排错和对比实验时不会因为数据划分不同而产生干扰。3. 模型选型与训练参数为什么小样本场景下ResNet比ViT更稳3.1 岩石识别任务的特点类间差异小、样本量小、纹理主导岩石识别的核心难点是类间差异非常细微。花岗岩和闪长岩在外观上可能只有石英含量和暗色矿物比例的差异碳酸盐岩的不同类型在偏光显微镜下才能区分。这种“细微纹理差异”决定了模型结构的选择方向需要足够多的低层特征提取能力同时对全局语义信息的需求反而不像做自然图像分类那么高。另外岩石数据集通常不大——公开的岩石图像数据集一般也就几千张图分为几类到十几类不等很多细分岩性的样本量不到100张。这种数据规模下从零训练一个深度网络基本是死路。常见做法是使用在ImageNet上预训练过的模型做迁移学习冻结大部分层只微调最后几层。ResNet18和ResNet50是我在岩石识别任务里最常用的两个骨干网络它们结构稳定、显存占用小、在中小数据集上不容易过拟合。ViT这种基于Transformer的模型在小数据集上表现反而不好因为它的归纳偏置弱需要大量数据才能学到有效的特征表达——这是个典型的“大力出奇迹”模型样本量不够时精度反而会被ResNet吊打。3.2 用PyTorch搭一个带迁移学习的训练脚本ResNet18 微调策略下面给出一个可以完整跑通的训练脚本。它做的事情是加载预训练的ResNet18、替换最后一层全连接分类头、冻结大部分层只微调最后一两个Block、用AdamW优化器训练。import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms, models from torch.utils.data import DataLoader # 数据增强岩石识别不适合做大幅旋转做小角度旋转和小尺度裁剪即可 train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rock_dataset_split/train, transformtrain_transform) val_dataset datasets.ImageFolder(rock_dataset_split/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers2) # 加载预训练模型替换分类头 model models.resnet18(pretrainedTrue) num_features model.fc.in_features num_classes len(train_dataset.classes) model.fc nn.Linear(num_features, num_classes) # 冻结backbone只训练最后两层的残差块和分类头 for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False # 注意只给需要梯度的参数传入优化器否则冻结部分会白白计算 optimizer optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-4, weight_decay1e-4 ) criterion nn.CrossEntropyLoss() # 训练循环 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) for epoch in range(30): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 每轮在验证集上验证一次 if (epoch 1) % 5 0: model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() acc 100 * correct / total print(fEpoch {epoch1}, Loss: {running_loss:.4f}, Val Acc: {acc:.2f}%) else: print(fEpoch {epoch1}, Loss: {running_loss:.4f}) torch.save(model.state_dict(), rock_resnet18.pth)训练核心信息学习率1e-4配合AdamW权重衰减1e-4防止过拟合冻结策略是只训练layer4和fc层前面的ResNet层作为固定的特征提取器。RandomResizedCrop的scale(0.8, 1.0)限定裁剪范围不做过度缩放因为岩石纹理的正确尺度对识别很关键。RandomRotation(15)做小角度旋转增强不超过15度是因为大角度旋转会让岩石结构角度特征失真。ColorJitter调整亮度和对比度模拟不同光照环境下的拍摄差异。参数调整的关键点如果你的数据集特别小总样本500建议把训练轮数从30降到15左右因为大模型的验证精度通常在10轮之后就饱和了继续训练只会让验证集精度停滞、测试集精度下降如果样本量特别大比如有上万张岩石薄片扫描图可以尝试解冻整个网络把学习率降到1e-5做全局微调精度通常还能再涨2~3个百分点。3.3 类别不平衡的损失函数改造用Weighted CrossEntropyLoss前面提到岩石数据集的类别分布经常不均衡这时候直接使用朴素的CrossEntropyLoss会出问题模型为了降低整体损失会把所有样本预测为样本量最大的那个类因为这么做对Loss的降低幅度最大。这种模型在验证集上看起来精度很高——比如大类别占了70%的样本你全猜大类别也有70%的准确率——但对小类别的识别几乎为0。解决办法是给损失函数每个类别一个权重样本数少的类别权重高这样模型会主动纠正对小类别的错误预测。import torch.nn.functional as F class_counts [len(os.listdir(os.path.join(rock_dataset_split/train, c))) for c in train_dataset.classes] total sum(class_counts) weights [total / (len(class_counts) * c) for c in class_counts] weights_tensor torch.tensor(weights, dtypetorch.float32).to(device) criterion_weighted nn.CrossEntropyLoss(weightweights_tensor)这个加权策略的本质是让每个类别的Loss贡献在期望上接近相同。公式是每个类别的权重等于总样本数除以类别数和该类样本数的乘积。举个例子三类样本分别是100、50、10那么权重就是(160/(3×100), 160/(3×50), 160/(3×10)) (0.533, 1.067, 5.333)第三类的权重是第一类的十倍。效果就是模型每猜错一个小类别样本带来的惩罚是大类别样本出错的好几倍。用了加权损失之后小类别的识别精度通常会明显改善但大类别的精度可能会有1~2个百分点的轻微下降这是模型总体的平衡之举正常可接受。4. 训练过程中的参数调优与收敛判断如何不被Loss走势骗了4.1 学习率调度策略每10轮衰减一次比固定学习率好用训练时如果全程用一个固定的学习率通常会出现这种情况前10轮Loss快速下降后面陷入平台期Loss波动但就是不降。很多人误以为模型已经收敛了其实只是学习率太大在最优解附近来回震荡。常见的做法是用步进衰减每训练一定轮数后把学习率乘一个衰减系数。下面这段代码在每隔8轮把学习率乘以0.5配合前面命令保存的最佳模型权重。from torch.optim.lr_scheduler import StepLR scheduler StepLR(optimizer, step_size8, gamma0.5) for epoch in range(30): # ...训练循环和验证代码... scheduler.step() # 每轮更新学习率 current_lr optimizer.param_groups[0][lr] print(fEpoch {epoch1}, current_lr: {current_lr})实际使用中StepLR的step_size和gamma需要配合数据集大小调整。数据集越小收敛越快step_size应该设得更小。以ResNet18和224x224输入为例如果训练集只有300张图每轮只有十来个Batch模型在第4轮开始就进入平台期此时step_size6比较合适如果训练集有2000张step_size10更稳妥。判断依据看验证集Loss曲线如果连续2到3轮验证Loss没有下降就说明学习率需要衰减了。还有一种更省心的方式是用ReduceLROnPlateau它会根据验证Loss是否停止下降自动调低学习率不需要预先设置衰减步数适合没有经验的初学者使用。4.2 模型保存与加载机制不能只保存最后一轮权重一个典型的初学者误区是训练结束后直接torch.save(model.state_dict(), final.pth)保存的是最后一轮权重。但深度学习训练里最后一轮权重通常不是验证集上最好的那一版因为训练末期模型已经在验证集上过拟合了。正确做法是每轮在验证集上评估一次记录验证精度最高的那一轮权重把它单独保存下来。训练结束后再手动加载这份最优权重做测试集评估。best_acc 0.0 best_epoch -1 for epoch in range(30): # ...训练代码... model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc 100 * correct / total if val_acc best_acc: best_acc val_acc best_epoch epoch 1 torch.save(model.state_dict(), best_rock_model.pth) print(fEpoch {epoch1}: 新的最佳精度 {val_acc:.2f}%) print(f训练完成最佳模型在第{best_epoch}轮精度{best_acc:.2f}%)加载模型非常简单先实例化和训练时一样结构的模型然后load_state_dict即可。但注意如果你训练时用了model models.resnet18(pretrainedTrue)然后修改了fc层加载时也需要做完全相同的修改否则参数维度对不上报错。另外如果你在训练脚本里对类别标签做过映射比如把类别名字符串映射成了数字索引加载时需要保留下这份映射关系因为部署推理时你需要知道模型的输出索引对应哪个岩性。建议把这个映射表存成一个JSON文件随模型一起保存不然模型换台机器就跑不了。import json class_to_idx train_dataset.class_to_idx # torchvision返回的字典 idx_to_class {v: k for k, v in class_to_idx.items()} # 反转 with open(idx_to_class.json, w) as f: json.dump(idx_to_class, f, indent2)这个JSON文件是模型完整的一部分。很多人在本地训练精度很高换到别的机器上推理时直接翻车而且报错信息非常隐晦——模型推理输出了一堆概率但不知道每个概率对应什么岩性。保存了映射关系这类问题一劳永逸。4.3 用TensorBoard看Loss分布发现过拟合的第一现场我不建议只在终端打印Loss数字因为数字变化趋势很难直观感知。TensorBoard是PyTorch官方集成的可视化工具torch.utils.tensorboard模块直接可用。装上之后把每轮的训练Loss和验证Loss记录下来画在同一个坐标轴里过拟合的特征非常明显训练Loss持续下降但验证Loss在第10轮左右开始反弹甚至一路走高。看到这个曲线就该立刻停止训练或加大正则化强度。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/rock_experiment) for epoch in range(30): # ...训练代码拿到train_loss和val_acc后... writer.add_scalar(Loss/train, running_loss, epoch 1) writer.add_scalar(Accuracy/val, val_acc, epoch 1) writer.add_scalar(LR, current_lr, epoch 1) writer.close()TensorBoard曲线还有一个作用判断模型是否太简单学习不到特征。如果训练Loss和验证Loss在同一个比较高的位置双双持平不下降比如Loss始终在1.5左右震荡说明模型容量不够或者数据本身太难学这时候要考虑换更大的模型ResNet50替代ResNet18或检查数据标注是否正确。而如果训练Loss降到很低但验证Loss明显偏高说明模型记住训练集了优先降低训练轮数或增加Dropout/Weight Decay而不是加大数据增强——这两个问题的解法方向完全不同没有曲线图做依据很容易走错路。5. 推理部署与精度评估模型训练出来只是开始5.1 单张图片预测脚本从加载模型到输出岩性训练完成之后模型还是一个.pth权重文件你需要写推理脚本把它包装成能对单张图片预测的函数。这里有个很容易踩的坑训练时你用了Normalize和Resize等预处理推理时也必须用一模一样的预处理否则输入分布不同模型输出结果会失真甚至完全错乱。下面这份推理代码严格按照训练时的val_transform做预处理输出每个类别的概率分布。import torch from torchvision import transforms, models from PIL import Image import json device torch.device(cuda if torch.cuda.is_available() else cpu) # 重建模型结构 model models.resnet18(pretrainedFalse) num_classes 10 # 替换成你自己数据集的实际类别数 model.fc torch.nn.Linear(model.fc.in_features, num_classes) model.load_state_dict(torch.load(best_rock_model.pth, map_locationdevice)) model.to(device) model.eval() # 加载类别映射表 with open(idx_to_class.json, r) as f: idx_to_class json.load(f) # 与训练val_transform保持一致 transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def predict_rock(img_path): img Image.open(img_path).convert(RGB) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) probabilities torch.softmax(outputs, dim1)[0] top_prob, top_idx torch.max(probabilities, 0) class_name idx_to_class[str(top_idx.item())] # 输出Top-1和Top-3结果便于人工确认 top3 torch.topk(probabilities, k3) print(f预测结果: {class_name}, 置信度: {top_prob.item()*100:.1f}%) for i in range(3): idx top3.indices[i].item() print(f {idx_to_class[str(idx)]}: {top3.values[i].item()*100:.1f}%) return class_name, top_prob.item() predict_rock(test_images/granite_sample.jpg)这份代码暴露了推理的两个细节第一map_locationdevice要写在torch.load里否则把GPU训练的权重加载到CPU机器上会报显存相关的错误第二torch.softmax需要指定dim1对二维张量来说dim1是按行做归一化才能得到每个样本的概率分布。输出Top-3而不是只输出Top-1是因为岩石类间相似度高模型可能只给出略高于其他类别的首选答案看Top-3能帮你判断模型是否处于“犹豫”状态。5.2 模型精度评估体系准确率、混淆矩阵与每个类别的召回率单看准确率是深度学习里最危险的做法。岩石识别场景有个经典案例某模型整体准确率90%但实际只把样本量最大的两个类别分对了另外几个类别全错而业务方恰恰需要识别的是那些稀有岩性。所以完整的评估至少要有三张表整体准确率、每个类别的召回率和精确率、混淆矩阵。这三张表能告诉你模型对哪个类别最不敏感、哪两个类别经常互相混淆、稀有类别是否完全没被学到。import numpy as np from sklearn.metrics import confusion_matrix, classification_report def evaluate_model(model, dataloader): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) print(混淆矩阵:) print(cm) # 每个类别的精确率、召回率、F1 report classification_report( all_labels, all_preds, target_nameslist(idx_to_class.values()), digits4 ) print(report) return cm, report test_dataset datasets.ImageFolder(rock_dataset_split/test, transformval_transform) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse) cm, report evaluate_model(model, test_loader)我强调两个评估细节第一这个评估必须只在测试集上做一次不能在验证集上调参调完又用同一份数据评估说精度有多高那相当于考试前先看了答案——验证集上表现不错但在测试集上严重翻车这是模型过拟合验证集的结果。第二classification_report输出的每个类别指标比整体准确率重要得多如果一个类别召回率只有30%说明这类样本大量被误判成了其他类别后续要根据混淆矩阵里的具体错判方向在数据增强时补充对应的样本形态。5.3 导出模型为ONNX脱离PyTorch环境做推理如果你想把模型放到生产环境用C或Java调用继续用PyTorch的.pth文件不方便——目标服务器未必安装了Python环境或者部署团队只认标准格式。ONNX是深度学习的通用交换格式支持把PyTorch模型导出为.onnx文件再借助ONNX Runtime在任意语言里加载推理。导出脚本很短但有几个参数需要留意。import torch.onnx model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, rock_model.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} }, opset_version17 ) print(ONNX导出完成)导出时最容易遇到的坑是模型包含某些ONNX不支持的算子导致导出失败或推理结果不一致。ResNet18、ResNet50这类经典结构基本不会有这个问题但如果你用的是较新的模型结构比如引入了某些自定义注意力模块导出前先在官方算子文档里查一下。dynamic_axes参数把batch维度设为动态这样导出的模型既能输入单张图片也能一次处理一批图片做推理服务时非常方便。opset_version建议设高一些旧版本opset可能缺少新算子的映射规则。6. 岩石识别的进阶方向从单标签分类到多标签任务与部署优化6.1 多标签与细粒度分类当岩石同时具备多个属性时怎么办基础版本的岩石识别是单标签分类——一张图对应一种岩性。但实际业务里常见的是多标签场景一块岩石可能同时标注了“碳酸盐岩”“细粒”“层理构造”等多个属性或者需要判断“是否是目标岩性”这个二分类结果。处理方法通常是把模型的最后一层从输出类别数改成输出属性数每个属性用一个Sigmoid激活函数损失函数改成BCEWithLogitsLoss训练数据集的标签从单个整数改成多个0/1向量。这个改造在代码层面改动很小但数据标注的工作量会显著增加——每个样本需要人工标注多个属性而不是只打一个标签。细粒度分类则是指同类岩石下的细分类型判断比如“火山岩”大类下区分玄武岩、安山岩、流纹岩。这类任务对模型要求更高单靠ResNet18输出特征可能不够常见做法是加入注意力机制模块如SE Block或在多个尺度上提取特征。如果你把标题里的项目跑通之后还有余力这是最有业务价值的进阶方向——很多地质单位对“大类识别对但细分岩性识别不准”的问题最为头痛。6.2 模型推理加速批量推理与半精度如果当前项目的数据量和调用频率上来了你会发现推理速度成为瓶颈。两个最常见的优化手段批量推理和半精度。批量推理是把多张图片合成一个Tensor一起过模型充分利用GPU并行能力半精度是把模型权重从FP32变成FP16或BF16几乎不损失精度但推理速度翻倍。PyTorch在推理阶段开启半精度非常简单主线代码在torch.no_grad()基础上包一层torch.autocast()。with torch.inference_mode(): with torch.autocast(cuda, dtypetorch.float16): outputs model(images)这段代码必须放在GPU环境下运行CPU不支持float16加速反而可能变慢。另外半精度推理的数值范围有限如果模型输出层接近1的概率被截断或者出现NaN说明网络的动态范围太大不适合直接用半精度可以退回单精度或在损失位点做保护。我做过的部署项目里ResNet18从FP32切换到FP16后推理延迟大约降低了35%精度波动通常在0.1%以内值得在生产环境尝试。6.3 最后的实战建议从“能跑通”到“能交付”的三道关卡我的个人经验是跑通一个标题里的项目很容易但交付成一个真正能用的系统还有三道关卡。第一道关卡是数据闭环——模型上线后要持续收集推理错误的样本并让人工复核然后把这些难例增量补充到训练集里否则模型在真实场景里永远只会越来越偏不会自我纠正。第二道关卡是边界认知——给模型设定置信度阈值当top1概率低于0.6时拒绝判定并转人工这比让模型硬猜然后给出错误答案更符合业务预期。0.6这个值可以根据混淆矩阵的结果调整目标是让“机器判错的概率”降到最低而不是让“机器判对的概率”最高。第三道关卡是复现保障——训练环境、数据集版本、依赖包版本都记录下来否则三个月后模型需要重训时你会发现环境完全无法复现。这三道关卡都用代码处理起来相对直接难的是在项目推进过程中坚持去做。做项目的过程中我犯过的最大错误是在模型精度上追逐最后一两个百分点而忽略了样本边界的定义——导致模型在实际场景里遇到训练分布外的岩石纹理时给出的高置信度预测完全是错的。现在我的做法是先检查模型对每类样本的置信度分布再优先解决低置信度区域的样本补充精度自然就上去了。希望这些经验对你做这个项目有实际帮助。本文还有配套的精品资源点击获取