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试一试。发现可以训练
