PIDNet（语义分割）排坑

最新推荐文章于 2025-03-17 17:58:34 发布

小马敲马

最新推荐文章于 2025-03-17 17:58:34 发布

阅读量1.4k

点赞数 24

分类专栏： paper阅读及复现文章标签： python 深度学习机器学习

本文为博主原创文章，未经博主允许不得转载。

本文链接：https://blog.youkuaiyun.com/buuliuda/article/details/143366297

版权

paper阅读及复现专栏收录该内容

6 篇文章

订阅专栏

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试一试。发现可以训练
在这里插入图片描述