如何对自定义数据集上的deeplabv3进行训练？

Question

我试着做图像分割，我了解了谷歌的工作，DeepLabv3。

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

陈，L.C.，Papandreou，G.，Schroff，F.和Adam，H.，2017。语义图像分割的四向卷积的再思考。arXiv预印版arXiv:1706.05587。

该体系结构用于对Pascal 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

• 对模型进行培训和评价。