当前位置: 首页 > article >正文

PIDNet(语义分割)排坑

PIDNet训练自己的数据集

    • 1. 前言
    • 2. 准备工作
    • 3. 配置环境
    • 4. 排坑过程
        • 4.1.1 configs增加了VOC文件夹 并在里面写了yaml参数文件
        • 4.1.2 加载VOC格式数据集的类
        • 4.1.3 train.py调试

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


http://www.kler.cn/a/373675.html

相关文章:

  • C++学习笔记3——存储持续性、作用域和链接性
  • VScode + PlatformIO 了解
  • 分布式事务Seata-AT模式
  • Linux_04 Linux常用命令——tar
  • 数据结构-希尔排序(ShellSort)笔记
  • AES_ECB算法C++与Java相互加解密Demo
  • HarmonyOS生命周期
  • 基于局部近似的模型解释方法
  • 【数据结构】ArrayList的模拟实现--Java
  • android12属性设置
  • 使用 NCC 和 PKG 打包 Node.js 项目为可执行文件(Linux ,macOS,Windows)
  • 设计一个灵活的RPC架构
  • AI代币是什么?AI与Web3结合的未来方向在哪里?
  • Transformer-BiGRU多特征输入时间序列预测(Pytorch)
  • WSGI、uwsgi与uWSGI
  • 【深度学习】用LSTM写诗,生成式的方式写诗系列之一
  • 下一代「自动化测试框架」WebdriverIO
  • STM32--STM32 微控制器详解
  • unity3d————Mathf.Lerp() 函数详解
  • 从0开始深度学习(21)——读写数据和GPU
  • 【Nas】X-DOC:Mac mini 安装 ZeroTier 并替换 planet 实现内网穿透
  • 人工智能中的机器学习和模型评价
  • RNN在训练中存在的问题
  • 常见的机器学习模型汇总
  • C++ 复习记录(个人记录)
  • 基于Multisim的四位抢答器设计与仿真