告别Cityscapes：手把手教你将DDRNet.pytorch项目迁移到自己的小数据集（以512x512细胞图为例）

从Cityscapes到细胞图像：DDRNet.pytorch项目迁移实战指南

当我们需要将开源的语义分割模型应用到自己的专业领域时，往往会遇到一个共同难题：如何把基于大型通用数据集（如Cityscapes）训练好的模型，快速适配到我们的小规模专业数据集上？本文将以512x512的细胞图像分割为例，带你完整走过DDRNet.pytorch项目迁移的全流程。

1. 项目迁移前的准备工作

在开始代码修改之前，我们需要先理清几个关键概念。DDRNet最初是为城市场景分割设计的，而我们要处理的是显微镜下的细胞图像，这两者在数据特性上有显著差异：

• 图像尺寸：Cityscapes通常使用2048x1024的高分辨率图像，而我们的细胞图像是512x512
• 类别数量：城市场景可能有19-34个类别，而细胞图像可能只需要区分3-5种状态
• 数据量：Cityscapes包含约5000张精细标注图像，而我们可能只有几百张标注样本

数据准备 checklist：

1. 确保图像和标签一一对应
2. 标签图像应为8位灰度图，像素值代表类别ID
3. 建议按7:2:1的比例划分训练集、验证集和测试集
4. 为每个类别分配连续的整数ID（从0开始）

提示：可以使用OpenCV的cv2.imread()加载标签图像时指定flag=0来确保以灰度模式读取

2. 自定义数据集类的实现

DDRNet默认使用Cityscapes数据集类，我们需要创建自己的数据集类。在lib/datasets/目录下新建CellDataset.py：

import os import numpy as np from PIL import Image from torch.utils.data import Dataset class CellDataset(Dataset): def __init__(self, root, list_path, transform=None): self.root = root self.list_path = list_path self.transform = transform self.img_list = [] self.label_list = [] with open(list_path, 'r') as f: for line in f: items = line.strip().split() self.img_list.append(items[0]) if len(items) > 1: # 训练/验证集有标签 self.label_list.append(items[1]) # 细胞图像特有的均值和标准差 self.mean = [0.485, 0.456, 0.406] # 需根据实际数据计算 self.std = [0.229, 0.224, 0.225] # 需根据实际数据计算 # 类别权重（处理类别不平衡） self.class_weights = [1.0, 1.5, 2.0, 1.8] # 示例值，需实际计算 def __len__(self): return len(self.img_list) def __getitem__(self, idx): image = Image.open(os.path.join(self.root, self.img_list[idx])).convert('RGB') if len(self.label_list) > 0: # 训练/验证集 label = Image.open(os.path.join(self.root, self.label_list[idx])) label = np.array(label) sample = {'image': image, 'label': label} else: # 测试集 sample = {'image': image} if self.transform: sample = self.transform(sample) return sample

关键修改点说明：

别忘了在lib/datasets/__init__.py中注册新数据集类：

from .CellDataset import CellDataset

3. 配置文件的关键参数调整

DDRNet使用YAML文件管理配置，我们需要修改experiments/cityscapes/ddrnet23_slim.yaml：

DATASET: NAME: 'cell' # 改为你的数据集名称 ROOT: './data/cell' # 数据根目录 NUM_CLASSES: 4 # 你的类别数 BASE_SIZE: 512 # 基础尺寸（原图大小） CROP_SIZE: 512 # 裁剪尺寸 TRAIN: BATCH_SIZE_PER_GPU: 4 # 根据GPU显存调整 LR: 0.01 # 小数据集可能需要更小的学习率 EPOCHS: 200 # 小数据集可能需要更多epoch AUG: SCALES: [0.5, 0.75, 1.0, 1.25, 1.5] # 缩放范围调整 FLIP: True # 水平翻转对细胞图像通常有效

小数据集训练策略调整建议：

• 学习率：初始值可设为0.01，配合学习率衰减策略
• 批量大小：在显存允许范围内尽可能大（但小批量可能导致训练不稳定）
• 数据增强：对细胞图像有效的增强包括：
  • 随机旋转（0-360度）
  • 颜色抖动（轻微调整对比度和亮度）
  • 弹性变形（模拟细胞形态变化）
• 早停机制：监控验证集指标，防止过拟合

4. 模型架构的适配修改

DDRNet的最后一层需要调整以匹配我们的类别数。修改lib/models/ddrnet_23_slim.py：

class DualResNet_imagenet(nn.Module): def __init__(self, block, layers, num_classes=4, planes=32, spp_planes=128, head_planes=128): super(DualResNet_imagenet, self).__init__() # ... 保持其他部分不变 ... self.final_layer = nn.Sequential( nn.Conv2d(planes * 4, head_planes, kernel_size=3, padding=1, bias=False), BatchNorm2d(head_planes), nn.ReLU(inplace=True), nn.Conv2d(head_planes, num_classes, kernel_size=1) # 输出通道改为你的类别数 )

对于小数据集，还可以考虑以下模型调整：

1. 减少通道数：将planes参数从32减小到24或16
2. 简化注意力机制：修改或移除部分DAPPM模块
3. 添加正则化：在关键位置增加Dropout层

5. 训练与评估技巧

5.1 小数据集训练策略

python train.py --dataset cell --cfg experiments/cell/ddrnet23_slim.yaml \ --batch-size 8 --lr 0.01 --epochs 200 --weight-decay 1e-4

小数据集训练的关键点：

• 迁移学习：加载Cityscapes预训练权重（除最后一层外）

  def load_pretrained(model, pretrained_path): pretrained_dict = torch.load(pretrained_path) model_dict = model.state_dict() # 过滤掉最后一层权重 pretrained_dict = {k: v for k, v in pretrained_dict.items() if not k.startswith('final_layer')} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)

• 类别平衡：在损失函数中使用类别权重

  criterion = nn.CrossEntropyLoss( weight=torch.tensor(dataset.class_weights).cuda() )

• 混合精度训练：减少显存占用，允许更大的batch size

  from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

5.2 评估与结果可视化

修改eval.py以保存预测结果：

def main(): # ... 原有代码 ... sv_pred = True # 确保开启结果保存 if sv_pred: sv_dir = os.path.join(output_dir, 'test_results') if not os.path.exists(sv_dir): os.mkdir(sv_dir) # 在预测循环中添加保存逻辑 for i, (image, label, filename) in enumerate(test_loader): # ... 预测代码 ... if sv_pred: pred = pred.cpu().numpy() for j in range(pred.shape[0]): save_pred(pred[j], sv_dir, filename[j])

结果分析表格：

针对上述问题，可能的改进措施：

1. 数据增强：增加更多样的细胞形态变化
2. 损失函数：对稀有类别（如坏细胞）增加权重
3. 模型简化：减少参数量防止过拟合
4. 伪标签：对未标注数据生成伪标签进行半监督学习

6. 实际应用中的优化技巧

在将模型部署到实际细胞分析流程中时，我们发现几个实用技巧：

技巧1：动态阈值调整

def postprocess(pred, class_thresholds=[0.5, 0.6, 0.7, 0.5]): # pred: [C,H,W]的预测logits probs = torch.softmax(pred, dim=0) final_mask = torch.zeros_like(pred[0]) for i, thresh in enumerate(class_thresholds): final_mask[probs[i] > thresh] = i return final_mask

技巧2：多尺度测试增强

def multi_scale_test(model, image, scales=[0.5, 0.75, 1.0, 1.25]): h, w = image.size final_pred = torch.zeros((num_classes, h, w)) for scale in scales: scaled_img = F.resize(image, (int(h*scale), int(w*scale))) pred = model(scaled_img.unsqueeze(0)) pred = F.resize(pred, (h, w)) final_pred += pred.squeeze(0) return final_pred / len(scales)

技巧3：模型量化部署

# 训练后量化 model = torch.quantization.quantize_dynamic( model, {nn.Conv2d, nn.Linear}, dtype=torch.qint8 ) torch.jit.save(torch.jit.script(model), 'ddrnet_cell_quantized.pt')