Training Fully Convolutional Network with VOC-style Custom Data

经典语义分割模型及数据集。

# 序

相较目标检测而言，语义分割进一步识别每一个像素的标号。该帖将利用 PyTorch 提供的\(\textbf{fcn_resnet50}\)训练\(\text{VOC2012}\)数据集。其中，损失函数为\(\textbf{交叉熵损失函数}\)，优化器选择\(\text{Adam}\)，数据增广使用到的类为\(\text{transforms.v2}\)和\(\text{transforms.tv_tensors}\)。

# 数据来源与预处理

数据来源

本示例使用的是\(2012\)年\(\text{PASCAL}\)的\(\text{Visual Object Classes (VOC)}\)竞赛数据集，可以通过该直链下载。\(\text{VOC2012}\)文件夹内文件树状结构如下。

├── Annotations
│   ├── 2007_000027.xml
│   ├── 2007_000032.xml
│   ├── 2007_000033.xml
│   └── ...
├── ImageSets
│   ├── Action
│   │   ├── jumping_train.txt
│   │   ├── jumping_trainval.txt
│   │   ├── jumping_val.txt
│   │   └── ...
│   ├── Layout
│   │   ├── train.txt
│   │   ├── trainval.txt
│   │   └── val.txt
│   ├── Main
│   │   ├── aeroplane_train.txt
│   │   ├── aeroplane_trainval.txt
│   │   ├── aeroplane_val.txt
│   │   └── ...
│   └── Segmentation
│       ├── train.txt
│       ├── trainval.txt
│       └── val.txt
├── JPEGImages
│   ├── 2007_000027.jpg
│   ├── 2007_000032.jpg
│   ├── 2007_000033.jpg
│   └── ...
├── SegmentationClass
│   ├── 2007_000032.png
│   ├── 2007_000033.png
│   ├── 2007_000039.png
│   └── ...
└── SegmentationObject
    ├── 2007_000032.png
    ├── 2007_000033.png
    ├── 2007_000039.png
    └── ...

其中：语义分割任务仅关注\(\text{ImageSets/Segmentation}\)、\(\text{JPEGImages}\)和\(\text{SegmentationClass}\)内的文件。文本文件记录训练样本文件名称； .jpg 文件为特征数据， .png 文件为标签数据，二者均为栅格数据。此外，栅格图像高宽不一，因此需要对他们进行\(\text{RandomResizedCrop}\)，这样才能组成批量。

构造训练图像数据迭代器

以便每次迭代时仅读入所需小批量数据，进而节省内存资源消耗。\(\text{VOC-style}\)语义分割数据集中特征与标签都使用栅格图像存储。图片中每个像素使用不同的颜色来表示不同的标签。因此迭代器不仅需要返回栅格样式的特征和标签数据，在处理图像时还需要一个\(\textbf{colormap}\)来将颜色映射至标签类别。这是大部分语义分割数据集的逻辑。

构造\(\textbf{colormap}\)

VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]

VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']


def voc_colormap2label():
    """Build the mapping from RGB to class indices for VOC labels."""
    colormap2label = torch.zeros(256 ** 3, dtype=torch.long)
    for i, colormap in enumerate(VOC_COLORMAP):
        colormap2label[
            (colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    return colormap2label

def voc_label_indices(colormap, colormap2label):
    """Map any RGB values in VOC labels to their class indices."""
    colormap = colormap.permute(1, 2, 0).numpy().astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
           + colormap[:, :, 2])
    return colormap2label[idx]

VOC_COLORMAP 列表中的\(\text{RGB}\)与 VOC_CLASSES 列表中的标签类别一一对应。上述函数将标签栅格图像文件\(\textbf{(3通道×长×宽)}\)的张量转为\(\textbf{(长×宽)}\)每个元素为目标类别索引的张量，转变后没有通道维度。

构造迭代器

class custom_VOCSegDataset(data.Dataset):
    def __init__(self, df, image_dir, transform=None):
        self.data_frame = df
        self.image_dir = image_dir
        self.transform = transform
        self.colormap2label = voc_colormap2label()

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.to_list()
        feature_img_name = os.path.join(self.image_dir, 'JPEGImages', f'{self.data_frame.iloc[idx, 0]}.jpg')
        feature_image = torchvision.io.read_image(feature_img_name, mode=torchvision.io.image.ImageReadMode.RGB)  # image = Image.open(img_name)
        label_img_name = os.path.join(self.image_dir, 'SegmentationClass', f'{self.data_frame.iloc[idx, 0]}.png')
        label_image = torchvision.io.read_image(label_img_name, mode=torchvision.io.image.ImageReadMode.RGB)
        if self.transform:
            feature_image, label_image = self.transform(feature_image, tv_tensors.Mask(label_image))
        return feature_image.type(torch.float32), voc_label_indices(label_image, self.colormap2label)

这里使用\(\text{torchvision.io}\)提供的方法替代 PIL 读取图片数据。此外，由于特征与标签均为栅格数据，需要使用\(\text{transforms.v2}\)同时增广这\(\textbf{两笔}\)数据并将其转为张量。需要注意的是，为确保如愿以偿对第二个传入参数\(\text{ (标签)}\)进行增广处理，这里先将其转为了\(\text{tv_tensors.Mask}\)。

不涉及批量大小，每一次返回的\(\text{标签栅格数据}\)不含通道维。此时不需要像SSD目标检测那样将通道维通过\(\text{unsqueeze(0)}\)插入回去，原因为\(\text{CrossEntropyLoss}\)设置了 reduction=’none’ ，详见此处。

读取并转换\(\text{train.txt}\)

trans = v2.Compose([
    v2.RandomResizedCrop(size=(320, 480)),
    v2.ToTensor()
])
voc_dir = f'/VOCdevkit/VOC2012/'
train_df = pandas.read_csv('/VOCdevkit/VOC2012/ImageSets/Segmentation/train.txt', skip_blank_lines=True, header=None, names=['img_name'])
train_dataset = custom_VOCSegDataset(df=train_df, image_dir=voc_dir, transform=trans)
train_iter = data.DataLoader(train_dataset, batch_size=16, shuffle=True)

# 训练准备

实例化模型

实例化\(\text{PyTorch}\)提供的\(\text{fcn_resnet50}\)，并采用预训练好的\(\textbf{backbone}\)。需要注意的是，\(\text{RGB(0,0,0)}\)代表背景标号，并且其索引为\(0\)。参数 num_classes 设为\(21\)则已经包括背景类别。这不同于目标检测边缘框的标签不含背景类的情况。

pyt_fcn = torchvision.models.segmentation.fcn_resnet50(weights_backbone=ResNet50_Weights.IMAGENET1K_V2, num_classes=21)
pyt_fcn.to('cuda:0')

定义损失函数和优化器

# 损失函数
loss = nn.CrossEntropyLoss(reduction='none')
# 优化器
optimizer = optim.Adam(pyt_fcn.parameters(), lr=0.0001, weight_decay=5e-4)

因为这里的\(\text{CrossEntropyLoss}\)设置了 reduction=’none’ 。这使得计算损失时，形状为\((21\times 320\times 480)\)的预测结果的每一个通道都与形状为\((320\times 480)\)的标签栅格数据计算损失。最后再使用 sum() 求和。

接下来使用\(\text{Adam}\)优化器基于\(\text{VOC2012}\)数据训练\(\text{fcn_resnet50}\)。

# 训练

定义参数

num_epoch = 20
best_model_loss = None  # 初始化最佳损失
valid_loss_list = []  # 记录所有epoch的验证集损失
valid_loss = torch.Tensor().to(device='cuda:0')  # 记录一个epoch内每次迭代的验证集batch损失。用于求平均，得到每个epoch的验证集损失

开始训练

每轮\(\text{epoch}\)都以学习到的现有参数进行一次评估，若得到更好的结果则保存模型。

for epoch in range(num_epoch):
    epoch_loss = 0
    pyt_fcn.train()
    for train_imgs, label_imgs in train_iter:
        optimizer.zero_grad()
        train_imgs, label_imgs = train_imgs.to('cuda:0'), label_imgs.to('cuda:0')
        l = loss(pyt_fcn(train_imgs)['out'], label_imgs).mean(1).mean(1).sum()
        l.backward()
        optimizer.step()
        epoch_loss += l.cpu().detach().numpy()
        print(l.cpu().detach().numpy())
    print(f"Epoch [{epoch + 1}/{num_epoch}], Loss: {epoch_loss:.4f}")
    pyt_fcn.eval()  # 转变为评估形态
    valid_loss = torch.Tensor().to(device='cuda:0')
    for valid_feature_imgs, label_imgs in valid_iter:
        valid_feature_imgs, label_imgs = valid_feature_imgs.to('cuda:0'), label_imgs.to('cuda:0')
        with torch.no_grad():
            loss_valid = loss(pyt_fcn(valid_feature_imgs)['out'], label_imgs).mean(1).mean(1).sum()
            valid_loss = torch.cat([valid_loss, loss_valid.reshape(-1)])
    valid_loss_list.append(valid_loss.to(device="cpu").numpy().flatten().mean())
    if best_model_loss is None:
        best_model_loss = valid_loss_list[-1]
        torch.save(pyt_fcn.state_dict(), '/VOCdevkit/fcn_best_model.ckpt')
        print('到达世界最高城——理塘!')
    if valid_loss_list[-1] < best_model_loss:
        best_model_loss = valid_loss_list[-1]
        torch.save(pyt_fcn.state_dict(), '/VOCdevkit/fcn_best_model.ckpt')
        print('到达世界最高城——理塘!')
    print(f'epoch_{epoch}的验证精度为{valid_loss.to(device="cpu").numpy().flatten().mean():.3f}')

# 推理

设置模型

\(16\text{G}\)显存\(16\)的\(batchsize\)不需要使用半精度就能使用该模型做推理。

pyt_fcn = torchvision.models.segmentation.fcn_resnet50(weights_backbone=ResNet50_Weights.IMAGENET1K_V2, num_classes=21)
pyt_fcn.load_state_dict(torch.load('/VOCdevkit/fcn_best_model.ckpt'))
pyt_fcn.to('cuda:0')
pyt_fcn.eval()  # 转变为评估形态

实例化验证集迭代器并得到\(16\)例验证集样本

valid_df = pandas.read_csv('/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt', skip_blank_lines=True, header=None, names=['img_name'])
valid_dataset = custom_VOCSegDataset(df=valid_df, image_dir=voc_dir, transform=trans)
valid_iter = data.DataLoader(valid_dataset, batch_size=16)
feature_imgs, label_imgs = next(iter(valid_iter))
feature_imgs = feature_imgs.to('cuda:0')

进行推理

output_img = pyt_fcn(feature_imgs)['out']

# 绘图

初始化\(\text{Figure}\)和\(\text{Axis}\)

fig, ax = pyplot.subplots(3, 8, figsize=(16, 6))

根据每一列\(\text{(共八列)}\)来绘制坐标系

for nc in range(8):
    ax[0, nc].imshow(feature_imgs[nc].permute(1, 2, 0).type(torch.int16).to('cpu'))
    ax[1, nc].imshow(label_imgs[nc].unsqueeze(0).permute(1, 2, 0).type(torch.int16).to('cpu'))
    ax[2, nc].imshow(output_img[nc].argmax(dim=0).unsqueeze(0).permute(1, 2, 0).to('cpu'))
pyplot.tight_layout()
pyplot.show()