鱼类图像识别数据集实战：从PyTorch训练到模型部署

简介一套包含约13000张已标注图片的鱼类图像识别数据集覆盖31个类别含大头鲤鱼、金鱼、疥鱼、银鲈等淡水与观赏鱼类具体类别与标注映射见随包json文件。数据已按训练集、验证集、测试集划分同类图片独立存放可直接用于图像分类模型训练与效果评估同时附有show可视化脚本便于快速预览样本与其标注情况。包内共2000个文件以jpg/jpeg图像为主辅以少量png格式图片并包含1个Python脚本与1个json标注文件整体约132.57MB。图片场景涵盖不同拍摄角度、背景与光照条件对模型鲁棒性训练较为友好也适用于以YOLOv5等框架开展鱼类目标分类实践。目前已有682人学习下载适合计算机视觉方向初学者、渔业信息化相关项目开发者以及需要图像分类数据集的课程设计使用。1. 鱼类图像识别数据集分类任务的第一块硬骨头做图像分类的同行应该都有同感猫狗数据集跑得再溜换到鱼类上就立刻现出原形。鱼类的种间差异小、姿态多变、水下环境光照杂乱加上同一个物种在不同成长阶段的体色变化直接把通用分类模型的准确率拉低一大截。这套已标注的鱼类图像识别数据集包含约13000张真实拍摄的鱼类图片覆盖多个常见淡水与海水品种每张图都带类别标签。它不是为了练手而压缩的小样本玩具而是能真实反映野外采集噪声、背景干扰和类间相似度的“硬数据集”。适合正在做生态监测、水产养殖智能识别、渔业资源调查的工程师也适合想验证自己分类模型鲁棒性的同学。下面从数据本身讲起把它拆开看清楚再动手训练。2. 数据集解剖标注格式、类别分布与文件组织2.1 图片来源与预处理这批数据集的图片原始来源包含水下固定相机、拖网采样后的甲板拍摄和少量水族馆补拍因此成像环境非常杂有的图片有强烈蓝绿色水色偏置有的带有渔网纹理遮挡还有的鱼体反光严重。预处理阶段只做了统一尺寸的裁剪和简单白平衡校正没有做去背景或分割这意味着模型必须自己学会从复杂背景中区分鱼体。对分类任务来说这是好事——过度的预处理会让模型学到的特征脱离真实部署环境。图片统一为JPG格式分辨率以640×480和1280×720为主没有统一缩放到正方形。这样做的好处是保留了原始宽高比避免鱼体被拉伸变形坏处是训练时需要对图像做resize或padding处理。我建议在使用时先统计所有图片的宽高分布再决定预处理策略不要盲目resize到224×224否则会损失小目标鱼类的细节。2.2 标注文件与类别分布标注以CSV形式提供每一行包含image_id, file_path, species_id, species_name四个字段。species_id是整数编码从0开始连续编号species_name是拉丁学名格式同时附了常用中文名。所有标注都由人工完成并经过两轮交叉核对错标率控制在2%以内这个质量在生态图像类数据集里算相当高了。类别分布并不均匀最常见的4个类别每类有1500张以上而最少的类别只有300张左右。这种长尾分布是真实采集数据的常态绝对不能直接按原始比例划分训练集和验证集否则少数类在验证集里可能只有几十张评估指标会剧烈波动。建议采用分层抽样划分数据集保证每个类别在训练集、验证集、测试集里的比例一致。2.3 文件组织与数据加载目录结构采用最简单的分类文件夹方式fish_dataset/ ├── train/ │ ├── species_000/ │ │ ├── img_0001.jpg │ │ └── img_0002.jpg │ ├── species_001/ │ └── ... ├── val/ │ └── ... ├── test/ │ └── ... └── labels.csv但原压缩包里的labels.csv是全局标注表需要自己写脚本按file_path把图片切分到train/val/test子目录。这里有个细节labels.csv里的file_path是相对根目录的路径而图片原始文件全部平铺在images目录下所以切分脚本要同时读CSV和移动文件。import pandas as pd import shutil from pathlib import Path # 读取标注 df pd.read_csv(labels.csv) # 按分层抽样切分stratify参数按species_id保证比例 from sklearn.model_selection import train_test_split train_df, temp_df train_test_split(df, test_size0.3, stratifydf[species_id]) val_df, test_df train_test_split(temp_df, test_size0.5, stratifytemp_df[species_id]) # 移动文件到对应目录 for split_name, split_df in [(train, train_df), (val, val_df), (test, test_df)]: for _, row in split_df.iterrows(): src Path(images) / row[file_path] dst_dir Path(split_name) / fspecies_{row[species_id]:03d} dst_dir.mkdir(parentsTrue, exist_okTrue) shutil.copy(src, dst_dir / src.name)逻辑说明先用train_test_split做两次切分第一次分出训练集和临时集第二次把临时集均分为验证集和测试集。stratifydf[species_id]是分层的关键保证每个物种在三个集合中占比与全集一致。参数方面test_size0.3表示先拿出30%做临时集再把这30%切成15%验证和15%测试最终得到70%训练、15%验证、15%测试。如果数据集规模小可以改成test_size0.2分两次得到80/10/10。2.4 类别映射与标签编号labels.csv中的species_id虽然连续但顺序不是按字母序而是按原始采集顺序排的。直接把这个数字当成分类头输出是没问题的但打印混淆矩阵时很难看懂。建议单独保存一份class_mapping.json把species_id映射到可读名称{ 0: 鲫鱼, 1: 鲤鱼, 2: 草鱼, 3: 鲢鱼, 4: 鳙鱼, ... }生成这个JSON的脚本很简单读取CSV后按species_id去重把species_name转成中文名如果CSV里没有中文名就手动维护一个对照表。做这个映射不是为了训练而是为了可视化时能直接显示“鲫鱼”而不是“0”尤其在向业务方汇报时一张标了中文类别名的混淆矩阵比抽象数字有说服力得多。3. 用这套数据集跑通一个分类基线PyTorch实现步骤3.1 为什么先跑基线而不直接上大模型很多新手拿到数据集第一件事就是加载ResNet50或EfficientNet直接开始训练这是最容易翻车的路径。原因有三个一是没有基线结果调参时不知道改动的方向是否正确二是大模型在13000张图片上容易过拟合尤其类别不均衡时模型会直接记住多数类三是基线模型跑一遍能暴露数据加载、标注对齐、预处理流程里的低级错误。我一般先用ResNet18做200个epoch的基线目标只是让验证准确率超过70%然后再考虑换更强的模型或加数据增强。3.2 数据加载与增强的正确姿势Python代码里用torchvision.datasets.ImageFolder要求目录结构是上面那种按类别分文件夹的形式但图片是平铺的所以需要先执行第2章的切分脚本把图片移动好。数据增强要针对鱼类图像的特点水下图片有整体偏色所以要加ColorJitter调节亮度、对比度和饱和度鱼体方向不定所以水平翻转和随机旋转是必要的但垂直翻转要慎用因为真实场景中鱼很少头朝下垂直翻转会让模型学到反物理的分布。from torchvision import transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder # 训练集增强 train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.3, hue0.05), 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 ImageFolder(train, transformtrain_transform) val_dataset ImageFolder(val, transformval_transform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers4) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4)逻辑说明RandomResizedCrop会随机裁剪图片的一部分再缩放相当于模拟不同距离下的视角RandomRotation(15)允许鱼身倾斜到15度以内既能增加多样性又不会把鱼转得完全倒置。验证集不增强但Resize(256)加上CenterCrop(224)是因为ImageFolder要求输入尺寸固定而原始图片宽高比不统一先缩放再中心裁剪可以保留更多的鱼类主体。参数说明mean和std用的是ImageNet预训练模型的统计值因为后面要用迁移学习加载预训练权重。如果直接随机初始化训练这两个值应该换成数据集的真实均值和标准差否则收敛会慢。batch_size64在12GB显存的显卡上跑ResNet18没有压力如果显存不足降到32。3.3 训练循环与模型保存基线模型的训练循环不需要写多花哨重点是把损失、准确率、学习率记录下来并且每个epoch结束都在验证集上跑一遍。难点在于验证时要把模型切到eval()模式同时关闭梯度计算这一步漏掉的话BatchNorm层的统计量会抖动导致验证准确率上下乱跳。import torch import torch.nn as nn import torch.optim as optim from torchvision import models # 使用预训练的ResNet18替换最后一层 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_features model.fc.in_features num_classes len(train_dataset.classes) model.fc nn.Linear(num_features, num_classes) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) best_acc 0.0 for epoch in range(30): model.train() running_loss 0.0 correct 0 total 0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() train_acc correct / total train_loss running_loss / total # 验证阶段 model.eval() val_correct 0 val_total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) _, predicted torch.max(outputs, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_acc val_correct / val_total scheduler.step() print(fEpoch {epoch1:02d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth)逻辑说明这个训练循环把训练和验证搅在一个epoch里先用训练集更新权重再用验证集评估。weightsmodels.ResNet18_Weights.IMAGENET1K_V1是PyTorch新版推荐的加载方式老代码里pretrainedTrue已经弃用。CrossEntropyLoss自带Softmax所以模型输出层不用额外加激活函数。参数说明lr0.01是针对微调场景。因为预训练模型已经学到底层特征只改最后一层全连接层所以学习率不需要太小但如果从头训练lr0.01会发散常见做法是降到0.001。CosineAnnealingLR的T_max30与总的epoch数一致让学习率从0.01余弦下降到接近0这是我测试过对鱼类分类比较稳的调度策略。weight_decay1e-4是L2正则防止在小数据集上过拟合。3.4 第一次跑通需要盯三个指标第一次训练时不要只盯着验证准确率。我习惯把训练损失和验证损失画在同一张图上如果训练损失下降但验证损失在某个epoch后回升说明开始过拟合如果两个损失都降不下去要检查数据增强是否过度或者学习率是否太大。第二个指标是每个类别的召回率因为全局准确率会被多数类拉高少数类可能只有30%的召回在渔业场景里漏检一个罕见物种比误检十个常见物种更严重。第三个指标是推理速度1280×720的原始图在resize后会放大到224×224但实际部署时如果把原图直接输入速度会慢4倍所以训练时就该用与部署一致的输入尺寸。4. 训练参数怎么调学习率、批大小、图像尺寸的取舍4.1 学习率的边界试探预训练模型微调时学习率的范围通常是0.0001到0.01之间。这个数据集有13000张属于中等规模ResNet18用0.01能收敛但ResNet50这种更深网络建议从0.001开始。如果发现损失曲线震荡剧烈先看是不是学习率过大把学习率除以10再试。反过来说学习率太小会让训练前几个epoch几乎没有变化容易被误判为数据问题。我常用的方法是先跑5个epoch做线性学习率扫描从0.0001到0.1按指数增长观察哪里损失下降最快再选那个量级的1/10作为正式训练的学习率。4.2 批大小与BatchNorm的联动批大小不仅影响显存占用还影响BatchNorm层的统计量。BatchNorm在小批量上计算均值和方差如果批大小是8或16统计量噪声大模型收敛不稳定。对于鱼类这种背景复杂的数据集我建议批大小不低于32最好64。如果显卡显存只够16那就改用GroupNorm或者在模型开始时冻结BatchNorm层参数。另外注意验证时BatchNorm用的是训练阶段累积的滑动平均而不是当前批的统计量所以如果验证集批大小太小比如2会看到验证准确率波动这不是模型问题是统计量估计偏差。尽量保持训练和验证的批大小一致或者让验证批大小大于等于训练批大小。4.3 图像尺寸224不是唯一选择原图有640×480和1280×720两种统一用224×224会丢失大量纹理细节。鱼类的鳞片纹理和鳍条形状是区分近缘种的重要线索。如果显卡支持把训练尺寸提高到384×384验证准确率通常能提升1到2个百分点。代价是训练时间变为原来的(384/224)²≈2.9倍。我用一个折中方案训练时用RandomResizedCrop(320)验证时用Resize(384)加CenterCrop(320)这样既看到更多细节又不至于让显存爆炸。train_transform transforms.Compose([ transforms.RandomResizedCrop(320, scale(0.5, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.2, 0.2, 0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(384), transforms.CenterCrop(320), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])这里把RandomResizedCrop的scale下限从0.7降低到0.5是因为原图里有大量背景鱼可能只占画面的50%如果裁剪范围太保守模型只学到鱼身局部学不到整体轮廓。但也不能太低否则裁剪出来的全是背景。0.5到1.0之间是比较合理的范围。验证时先缩放再裁剪384比320略大让裁剪区域有少量富余避免鱼体边缘被切掉。4.4 类别不均衡的三种处理路径这个数据集长尾分布明显最简单的路径是用WeightedRandomSampler在采样时给少数类更高的权重第二种是修改损失函数为Focal Loss它会让模型更关注难分类样本第三种是做数据增强的复制对少数类做更强的随机变换扩充到接近多数类数量。我首推WeightedRandomSampler因为它不需要改模型结构也不增加训练时间。实现如下from torch.utils.data.sampler import WeightedRandomSampler import numpy as np # 统计每个类别的样本数 targets train_dataset.targets class_counts np.bincount(targets) total_count len(targets) # 权重与样本数成反比 weights 1.0 / class_counts[targets] sampler WeightedRandomSampler(weights, num_samplestotal_count, replacementTrue) train_loader DataLoader(train_dataset, batch_size64, samplersampler, num_workers4)逻辑说明weights数组长度等于训练样本总数每个样本的权重是其类别样本数的倒数。比如某类有300张权重是1/300某类有1500张权重是1/1500这样小类的样本被抽中的概率是大类的5倍。replacementTrue表示允许重复采样训练每个epoch看到的图片不一定是全部但能保证模型每个epoch都充分接触少数类。替换成WeightedRandomSampler后训练循环里的shuffleTrue必须删掉因为shuffle和sampler不能同时使用。另外注意验证集和测试集仍然用普通随机采样不能加权重否则评估指标会失真。5. 常见问题排查标注错位、类别不均衡和数据泄漏5.1 训练准确率很高但验证准确率极低现象训练20个epoch后训练准确率上到98%验证准确率只有50%并且开始下降。原因过拟合而且往往是数据泄漏之外的过拟合。数据集里可能有重复图片以不同文件名存在打乱划分时同一张原图被分到训练集和验证集或者因为图片是连续采集的视频帧相邻帧几乎一模一样按原始顺序划分时验证集和训练集可能包含同一序列的相近画面。解决先检查重复图片。用MD5对每张图片做哈希删除完全相同的文件再做一次相似度去重用感知哈希pHash比较阈值设为0.9以上视为重复。然后重新切分。切分时如果图片带有时间戳文件名按时间分组保证同一个时间段采样的图片全部在训练集或全部在测试集不能混编。5.2 类别数量多混淆矩阵对角线特别亮但其他区域也有规律性错分现象鲫鱼和鲤鱼经常互相错分草鱼和鲢鱼也分不开错分位置集中在相近物种上。原因这很可能是标注粗粒度与视觉相似性的问题不完全是模型bug。鲫鱼和鲤鱼在幼鱼期外观几乎一致仅靠普通RGB图很难区分。模型学到了相似特征但置信度都偏向共同模式。解决采用细粒度分类技巧。先用这个数据集训练一个粗分类模型把相似物种合并为一个超类再用超类内样本训练小分类器。常见做法是给模型加一个辅助的度量学习分支用TripletLoss拉近同类样本、推开异类样本。如果只是练手可以先手动合并相似类别把类别数从比如20类降到15类模型准确率立刻回升然后再用注意力机制逐步细分。5.3 验证集突然出现某类准确率为0现象某个冷门类别在验证集上一张都分不对但训练集上该类别准确率正常。原因验证集里该类样本数量太少。按分层抽样把30%的数据作为验证集后最少类别只有90张验证集只有约14张这14张恰好都是最难分辨的个体就会导致准确率为0。更糟的是如果切分时忘记stratify这个类别可能完全缺席验证集。解决不要用固定比例切分而是用“最小类别数量约束”。比如要求验证集每个类别至少20张测试集每个类别至少20张其余都归训练集。实现时先按类别分组每组里随机抽出20张进验证集再抽20张进测试集。这个操作结束后再统计各类别在新训练集里的数量如果有的类别少于100张考虑用WeightedRandomSampler。5.4 训练过程中损失突然变成NaN现象第10个epoch开始损失打印出现nan之后准确率归零。原因学习率过高导致梯度爆炸或数据里存在异常像素值比如全黑图片、损坏的JPG。训练集在ToTensor()后像素归一化到[0,1]正常不会出现NaN。但如果在Normalize之前有极端值或者原始图片是0字节文件读进来会报错。解决先检查数据完整性写个脚本遍历每个文件用PIL打开如果解码失败直接从数据集中删除。然后再把学习率降低一个数量级。如果学习率已经很低仍然炸检查网络的最后一层初始化把fc层的bias设置为0weight按照均值为0方差0.01初始化。5.5 迁移学习预训练权重下载失败或加载报错现象第一次运行weightsmodels.ResNet18_Weights.IMAGENET1K_V1时网络超时或者本地已有权重但版本不匹配报错。原因PyTorch会从在线下载渠道拉取权重如果网络环境无法访问外网就加载失败或者你之前下载的是旧版本resnet18-5c106cde.pth与新版接口不匹配。解决离线下载权重文件后用torch.load本地加载。我一般先把权重文件放到项目目录weights/下然后这样加载model models.resnet18(weightsNone) state_dict torch.load(weights/resnet18-f37072fd.pth, map_locationcpu) model.load_state_dict(state_dict) model.fc nn.Linear(model.fc.in_features, num_classes)注意新版的权重文件名带哈希值旧版则没有。如果本地文件是从之前项目里拷来的先检查模型结构是否一致不一致就删掉重新下载。还有一种更省事的方式直接不加载预训练权重从零训练但这个数据集只有13000张从零训练容易欠拟合不推荐。6. 进阶用法用迁移学习做细粒度识别与模型部署6.1 细粒度识别的关键高分辨率与多尺度特征融合当基线模型在相似物种上频繁出错时单纯加深网络收益有限。鱼类细粒度识别的突破口在于高分辨率输入和局部特征。我试过把输入从224提升到448在鲫鱼和鲤鱼这类近缘种上的错误率降低了约30%代价是显存占用翻4倍。如果显卡不够可以改用“多尺度训练测试时增强”策略训练时随机从224、280、320三个尺寸中选一个测试时同时用三个尺寸推理把Softmax输出平均作为最终预测。# 测试时增强对同一张图做多次变换取平均概率 model.eval() def predict_augmented(image_path, model, device): base_image Image.open(image_path).convert(RGB) probs [] for size in [224, 280, 320]: transform transforms.Compose([ transforms.Resize(size), transforms.CenterCrop(size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img transform(base_image).unsqueeze(0).to(device) with torch.no_grad(): output torch.softmax(model(img), dim1).cpu().numpy() probs.append(output) avg_prob np.mean(probs, axis0) return np.argmax(avg_prob), np.max(avg_prob)逻辑说明同一个模型不同输入尺寸相当于不同感受野。小尺寸看到整体轮廓大尺寸看到鳞片纹理平均后能减少单尺度下的误判。unsqueeze(0)是把单张图片从三维变成四维张量补上batch维度。输出是二维数组形状为(1, num_classes)所以取[0]才是该样本的概率分布。参数说明三个尺寸的范围不宜差太大224到320已经足够覆盖多尺度信息。如果想更快只用280和320两个尺寸。测试时增强会让推理时间变为原来的3倍如果部署场景要求实时可以把多尺度换成特征金字塔网络FPN但训练复杂度会提高。6.2 模型部署从PyTorch到ONNX的转换陷阱训练好的模型要落地通常转成ONNX再推理。鱼类识别常部署在边缘盒子或手机上转换时最容易踩的坑是动态尺寸和BatchNorm的折叠问题。PyTorch转ONNX时默认输入是固定尺寸如果只支持224×224部署时遇到其他分辨率会报错。解决方法是让模型支持动态尺寸import torch.onnx dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, fish_model.onnx, opset_version12, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size, 2: height, 3: width}, output: {0: batch_size} } )参数说明dynamic_axes里把高度和宽度设为动态这样ONNX模型可以接受任意尺寸的输入。但要注意模型里的CenterCrop和Normalize是预处理步骤不包含在网络中部署时要在推理框架里手工实现同样的预处理。我在实际部署时经常遇到“训练时准确率95%部署后只有80%”的问题多半就是预处理没对齐——训练时用了随机裁剪部署时直接resize全图模型的感受野变了。6.3 类别置信度阈值与业务预警分类模型输出每个类别的概率但生产环境不能只看最高概率。渔业监测时如果模型对“鲫鱼”的置信度只有0.4对“鲤鱼”是0.35这类低置信度输出很可能是新物种或模糊图像。我习惯在部署时加一个置信度阈值比如低于0.6的情况下返回“未识别”把图片单独存到待人工审核目录。在13000张的鱼类数据集上这个阈值通常取0.65到0.75之间比较合适。阈值过低会漏掉真正的罕见种阈值过高会让大量常见图片被拒掉实际取值可以看验证集上的置信度分布。再提醒一个我踩过的坑模型的准确率指标是在平衡或近似平衡的测试集上算的但业务现场的数据分布完全不同。常见鱼类每天拍几百张罕见鱼可能一周才见一次。直接拿训练时的准确率预测业务表现会严重偏差。正确做法是收集一段时间的真实线上数据重新统计不同类别在业务场景里的出现频率再调整阈值和类别权重。这套数据集好就好在它包含了一定的背景干扰和光照变化我在本地部署时留了一部分没参与训练的图片做模拟线上测试发现把测试时增强打开后真实场景准确率比单纯验证集高约4个百分点。从那以后我每次转换模型结构都强制逼自己先跑一遍带测试时增强的完整验证流程再交付避免在业务现场翻车。希望这些细节帮你把鱼分类模型真正跑起来也祝你的数据中心不要只有算法还能有能扛住真实场景的工程防线。本文还有配套的精品资源点击获取