如何在pytorch上的自定义数据集上训练deeplabv3

Question

我正在尝试进行图像分割，我了解了DeepLabv3在谷歌上的工作。

这是对论文的参考：https://arxiv.org/abs/1706.05587

陈，L.C.，帕潘德里欧，G.，施罗夫，F.和亚当，H.，2017。重新思考用于语义图像分割的atrous卷积。arXiv预印本arXiv:1706.05587。

该架构经过训练，可以对Pascal VOC 2012数据集的20+1个类(20个前台类和1个后台类(进行分割。

Pytorch在Pascal数据集上提供了预先训练的deeplabv3，我想在城市景观上训练相同的架构。因此，Pascal VOC数据集存在不同的类。我想知道做这件事的有效方法是什么？

目前，这是我写的唯一代码：

import torch
model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True)
model.eval()

Answer 1

• 编写自定义Dataloader类，该类应继承Dataset类并实现至少两个方法__len__和__getitem__
• 使用自定义数量的输出通道修改预训练的DeeplabV3头

from torchvision.models.segmentation.deeplabv3 import DeepLabHead
from torchvision.models.segmentation import deeplabv3_resnet101
def custom_DeepLabv3(out_channel):
model = deeplabv3_resnet101(pretrained=True, progress=True)
model.classifier = DeepLabHead(2048, out_channel)
#Set the model in training mode
model.train()
return model

• 训练并评估模型