PIDNet（语义分割）排坑

1. 前言

paper小修时reviewer说baseline太老，所以对CVPR2023的PIDNet进行复现，用于下游任务的baseline。所以写一个记录说明完整的过程，同时要说的是：
没有从0开始训练，用的PIDNet作者提供的在ImageNet上训练的预训练权重来继续训练的。数据集我是用的是无人机航拍数据集AeroScapes，VOC格式。解决了两个问题：
[1] 基础模型只支持cityscapes和camvid数据集，现在我支持了VOC格式的分割数据集。
[2] 基础模型是用double GPUs来训练 其中涉及到很多需要修改才能适配单张GPU训练

2. 准备工作

代码下载：去github官网下载，地址：PIDNet-code
论文下载：CVPR可以直接看PDF，地址：PIDNet-paper

3. 配置环境

这里就不一一介绍了 能看到这里 说明大家已经是炼丹老师傅了 一个个看缺啥库就下啥 最好用虚拟环境。

4. 排坑过程

4.1.1 configs增加了VOC文件夹 并在里面写了yaml参数文件

参考了作者的yaml文件写的
【1】更改了DATASET部分 因为我的数据集（aeroscapes）放在了data目录下，训练和验证的的索引也变成了我现在的路径。
【2】 MODEL部分 我增加了作者提供的ImageNet上的预训练权重的路径 pretrained_models/imagenet/PIDNet_S_ImageNet.pth.tar，这个权重是small版本，还有medium和large版本。small版本的权重下载地址是：PIDNet_S_ImageNet.pth.tar,下载后直接放到pretrained_models\imagenet\PIDNet_S_ImageNet.pth.tar路径下即可。
【3】其他超参数 微微动了点
以上修改后的完整代码如下：

# name: pidnet_vai_aero.yaml
CUDNN:
  BENCHMARK: true
  DETERMINISTIC: false
  ENABLED: true
GPUS: 0
OUTPUT_DIR: 'output'
LOG_DIR: 'log'
WORKERS: 3
PRINT_FREQ: 10

DATASET:
  DATASET: voc
  ROOT: 'data/'
  TEST_SET: 'data/aeroscapes/ImageSets/Segmentation/val.txt'
  TRAIN_SET: 'data/aeroscapes/ImageSets/Segmentation/train.txt'
  NUM_CLASSES: 11
MODEL:
  NAME: pidnet_small
  NUM_OUTPUTS: 2
  PRETRAINED: "pretrained_models/imagenet/PIDNet_S_ImageNet.pth.tar"
LOSS:
  USE_OHEM: true
  OHEMTHRES: 0.9
  OHEMKEEP: 131072
  BALANCE_WEIGHTS: [0.4, 1.0]
  SB_WEIGHTS: 1.0
TRAIN:
  IMAGE_SIZE:
  - 960
  - 720
  BASE_SIZE: 960
  BATCH_SIZE_PER_GPU: 6
  SHUFFLE: true
  BEGIN_EPOCH: 0
  END_EPOCH: 200
  RESUME: false
  OPTIMIZER: sgd
  LR: 0.005
  WD: 0.0005
  MOMENTUM: 0.9
  NESTEROV: false
  FLIP: true
  MULTI_SCALE: true
  IGNORE_LABEL: 255
  SCALE_FACTOR: 16
TEST:
  IMAGE_SIZE:
  - 960
  - 720
  BASE_SIZE: 960
  BATCH_SIZE_PER_GPU: 1
  FLIP_TEST: false
  MULTI_SCALE: false
  MODEL_FILE: ''
  OUTPUT_INDEX: 1

4.1.2 加载VOC格式数据集的类

打开datasets/init.py文件 加上

from .voc_dataloader import VOC as VOC

然后创建一个voc_dataloader.py，这个过程中仿照了作者cityscapes.py中的类，实现一样的初始化参数来匹配接口。这个py文件的代码如下：

import os

import cv2
import numpy as np
import torch
from PIL import Image

from .base_dataset import BaseDataset

class VOC(BaseDataset):
    # 数据集类的构造函数
    def __init__(self,
                 root,
                 list_path,
                 num_classes=11,
                 multi_scale=True, 
                 flip=True, 
                 ignore_label=255, 
                 base_size=2048, 
                 crop_size=(512, 1024),
                 scale_factor=16,
                 mean=[0.452, 0.502, 0.434],
                 std=[0.196, 0.161, 0.179],
                 bd_dilate_size=4):
        super(VOC, self).__init__(ignore_label, base_size,
                                  crop_size, scale_factor, mean, std)
        # 用构造函数的参数初始化类的成员变量
        self.root = root
        self.list_path = list_path
        self.num_classes = num_classes
        self.multi_scale = multi_scale
        self.flip = flip
        self.bd_dilate_size = bd_dilate_size
        self.ignore_label = ignore_label
        
        # 读取图像ID列表
        with open(self.list_path, 'r') as f:
            self.image_ids = [line.strip() for line in f.readlines()]

        self.files = []
        for image_id in self.image_ids:
            image_file = os.path.join(self.root, "aeroscapes/JPEGImages", image_id + '.jpg')
            label_file = os.path.join(self.root, "aeroscapes/SegmentationClass", image_id + '.png')
            self.files.append({
                "image": image_file,
                "label": label_file,
                "name": image_id
            })

        self.class_weights = torch.FloatTensor([0.80906685, 1.01004548, 
                                                1.15333424, 1.0154087,  
                                                1.20380376, 1.23027661,
                                                1.11751722, 0.98967911, 
                                                0.88035226, 0.79071721, 
                                                0.79979855]).cuda()

    def __len__(self):
        return len(self.files)
    
    
    def __getitem__(self, index):
        item = self.files[index]
        name = item["name"]
        image = cv2.imread(item["image"], cv2.IMREAD_COLOR)
        size = image.shape

        if 'test' in self.list_path:
            image = self.input_transform(image)
            image = image.transpose((2, 0, 1))

            return image.copy(), np.array(size), name

        label = cv2.imread(item["label"], cv2.IMREAD_GRAYSCALE)


        label[label == 255] = self.ignore_label
        label[label >= self.num_classes] = self.ignore_label
        label[label < 0] = self.ignore_label
        # label = torch.from_numpy(label).long()


        image, label, edge = self.gen_sample(image, label,
                                             self.multi_scale, self.flip, edge_size=self.bd_dilate_size)

        return image.copy(), label.copy(), edge.copy(), np.array(size), name

这里可以看到作者提前计算了数据集的mean，std和class_weights，所以我们也要计算出来并替换进去。根据我提供的代码来进行计算得到结果并替换进去，计算的代码是：

import os
import cv2
import numpy as np
from tqdm import tqdm

def compute_mean_std(image_paths):
    # 初始化
    channel_sum = np.zeros(3)
    channel_squared_sum = np.zeros(3)
    num_pixels = 0

    for img_path in tqdm(image_paths):
        img = cv2.imread(img_path)  # BGR 格式
        img = img / 255.0  # 归一化到 [0, 1]
        h, w, c = img.shape
        num_pixels += h * w

        # 累加每个通道的像素值
        channel_sum += np.sum(np.sum(img, axis=0), axis=0)

        # 累加每个通道的像素值的平方
        channel_squared_sum += np.sum(np.sum(np.square(img), axis=0), axis=0)

    # 计算均值
    mean = channel_sum / num_pixels

    # 计算标准差
    std = np.sqrt(channel_squared_sum / num_pixels - np.square(mean))

    return mean, std


def compute_class_weights(label_paths, num_classes, ignore_label=255):
    class_counts = np.zeros(num_classes)

    for label_path in tqdm(label_paths):
        label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
        # 忽略被标记为 ignore_label 的像素
        label = label[label != ignore_label]

        # 统计每个类别的像素数量
        for i in range(num_classes):
            class_counts[i] += np.sum(label == i)

    # 确保像素数量不为零，防止取对数和除零错误
    epsilon = 1e-6
    pixel_count = class_counts + epsilon

    # 调用您的函数计算类别权重
    class_weights = get_weight(num_classes, pixel_count)

    return class_weights

def get_weight(class_num, pixel_count):
    W = 1 / np.log(pixel_count)
    W = class_num * W / np.sum(W)
    return W



# 获取数据集中的所有图像路径
image_dir = 'data/aeroscapes/JPEGImages'
image_paths = [os.path.join(image_dir, filename) for filename in os.listdir(image_dir) if filename.endswith('.jpg')]

# 获取数据集中的所有标签路径
label_dir = 'data/aeroscapes/SegmentationClass'
label_paths = [os.path.join(label_dir, filename) for filename in os.listdir(label_dir) if filename.endswith('.png')]

mean, std = compute_mean_std(image_paths)
print("Dataset Mean: ", mean)
print("Dataset Std: ", std)
num_classes = 11  # 根据你的数据集设置
class_weights = compute_class_weights(label_paths, num_classes)
print("Class Weights: ", class_weights)

得到结果：

4.1.3 train.py调试

tools/train.py中的第一个函数是parse_args（），可以看到我们的参数配置文件现在是新的，所以把路径修改成现在的configs/VOC/pidnet_vai_aero.yaml，具体代码如下：

def parse_args():
    parser = argparse.ArgumentParser(description='Train segmentation network')
    
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        default="configs/VOC/pidnet_vai_aero.yaml",
                        type=str)
    parser.add_argument('--seed', type=int, default=304)    
    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()
    update_config(config, args)

    return args

调试main函数 给这个parse_args（）函数加上断点发现没有问题。step by step执行后发现，代码中有计算当前gpu数量的功能，代码如下：

gpus = list(config.GPUS)
    if torch.cuda.device_count() != len(gpus):
        print("The gpu numbers do not match!")
        return 0

由于我们现在只用一个GPU，所以去刚刚的pidnet_vai_aero.yaml文件中先把GPUS的参数从(0,1)改成0，回到train.py把这个计算gpu数量的代码修改成

gpus = config.GPUS
    if torch.cuda.device_count() != gpus+1:
        print("The gpu numbers do not match!")
        return 0

(tips:当然可以删除这一段代码。另外：由于现在GPUS的值是0，当中的list对于int形会报错 后面也要删除list。len测量这个list的长度也会报错，把后面出现的len（）函数都删了 )。
继续执行代码，发现batch_size作者是设置了一颗GPU是多少 乘以 GPU数量 直接删除len(gpus)即可。后面出现len(gpus)也记得删除，不然报错。

# 原本的代码
# batch_size = config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus)
# 现在的代码
batch_size = config.TRAIN.BATCH_SIZE_PER_GPU

继续执行发现刚刚的voc_dataloader没有给scale_factor传入参数 所以直接在voc_dataloader.py给它初始化为16。继续执行发现没有问题，在后面有一段代码

model = FullModel(model, sem_criterion, bd_criterion)
model = nn.DataParallel(model, device_ids=gpus).cuda()
#改成了
model = FullModel(model, sem_criterion, bd_criterion).cuda()

不用并行了 本人就用一张卡跑。测试没问题，继续执行。发现本地numpy版本新一些，继承了BaseDataset类，会出现np.int报错 所以把base_dataset.py中的np.int全部改成了int后错误消失。至此，成功run了。直接python tools/train.py试一试。发现可以训练