如何在 pytorch 上的自定义数据集上训练 deeplabv3？ [关闭]答案

【问题标题】：How to train deeplabv3 on custom dataset on pytorch? [closed]如何在 pytorch 上的自定义数据集上训练 deeplabv3？ [关闭]
【发布时间】：2021-01-01 15:09:26
【问题描述】：

我正在尝试做图像分割，我了解了 DeepLabv3 的 Google 工作。

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

Chen, L.C.、Papandreou, G.、Schroff, F. 和 Adam, H.，2017 年。重新思考用于语义图像分割的空洞卷积。 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()

【问题讨论】：

嗨@Muhammad Bilal，你看过这个了吗？ github.com/VainF/DeepLabV3Plus-Pytorch 它似乎完全按照你的要求做。

标签： python deep-learning pytorch deeplab

【解决方案1】：

编写custom Dataloader 类，该类应继承Dataset 类并实现至少2 个方法 __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

训练和评估模型。

【讨论】：