如何将35类城市景观数据集转换为19类?

时间:2019-06-18 13:23:19

标签: computer-vision pytorch image-segmentation

以下是我的代码的一小段。使用此工具,我可以在城市景观数据集上训练名为“ lolnet”的模型。但是数据集包含35个类别/标签[0-34]。

<form>
  <details>
    <input type="checkbox" required="required">
  </details>
  <input type="submit">
</form>

但是我只想在19个班级上训练我的模型。找到这19个类here。要训​​练的标签存储为“ ignoreInEval” = True。此数据集的pytorch Dataloader helper没有提供任何线索。

所以我的问题是如何使用pytorch的“ datasets.Cityscapes” api在此数据集的所需19个类上训练我的模型。

3 个答案:

答案 0 :(得分:1)

根据我的理解,您希望使用更少的课程从头开始训练新模型。有两种方法可以实现您的目标。通常首选的是第一种方式。

  • 用所有35个类训练模型,但在测试阶段(而不是评估阶段)忽略不需要的类。
  • 自定义Dataset类以仅从目标类加载图像。

首选第一种方法的原因是因为您可以训练更通用的模型,并且它在现实世界中应具有更好的性能。缺点是您需要更多的资源/时间来训练模型。

对于第二种方法,您需要创建一个自定义数据集类,如下所示。根据您的任务,准备一个列表,该列表指示训练和目标图像的路径以及标签。

from torch.utils.data.dataset import Dataset
from torchvision import transforms

class MyCustomDataset(Dataset):
    def __init__(self, ..., transforms=None):
        # custom list of image path, labels...
        self.transforms = transforms

    def __getitem__(self, index):
        # stuff
        ...
        data = # Some data read from a file or image
        if self.transforms is not None:
            data = self.transforms(data)
        # If the transform variable is not empty
        # then it applies the operations in the transforms with the order that it is created.
        return (img, label)

    def __len__(self):
        return count # of how many data(images?) you have

答案 1 :(得分:0)

您下载模型和权重。

import torch
import torch.nn as nn
import torchvision.models as models

r = models.resnet50(pretrained=True)

请注意,原始重寄有1000个类别/类别。因此,当您下载预训练的模型时,最后fc将用于1000个课程。

这是您拥有的forward()方法,并且在代码上方是您的模型。

您可以从原始的resnet50模型中删除最后一个fc完全连接的层,并添加具有完全19个类(19个输出)的新fc,并且可以仅针对最后一个层训练分类器。除最后一层外,其他各层均应冻结。

所以您将只学习所需的19个课程。


请注意,重新发送的__init__方法也可以使用类数,因此您可以尝试这样做,但是在这种情况下,您无法加载预训练的权重,因此您需要使用pretrained=False并且需要训练从头开始。

import torch
import torch.nn as nn
import torchvision.models as models

r = models.resnet50(num_classes=19, pretrained=False)

答案 2 :(得分:0)

已经有一段时间了,但是留下一个可能对其他人有用的答案

首先创建到19个类+背景的映射。背景与here所说的具有忽略标志的不太重要的类有关。

# Mapping of ignore categories and valid ones (numbered from 1-19)
    mapping_20 = { 
        0: 0,
        1: 0,
        2: 0,
        3: 0,
        4: 0,
        5: 0,
        6: 0,
        7: 1,
        8: 2,
        9: 0,
        10: 0,
        11: 3,
        12: 4,
        13: 5,
        14: 0,
        15: 0,
        16: 0,
        17: 6,
        18: 0,
        19: 7,
        20: 8,
        21: 9,
        22: 10,
        23: 11,
        24: 12,
        25: 13,
        26: 14,
        27: 15,
        28: 16,
        29: 0,
        30: 0,
        31: 17,
        32: 18,
        33: 19,
        -1: 0
    }

然后为您加载用于训练的每个标签图像(每个像素包含一个类的灰度图像,其灰度模式为“ {city} __ {number} _ {number} _gtFine_labelIds.png“),请在下面运行功能。 / p>

它将根据上面的映射转换每个像素,并且标签图像(蒙版)现在将只有20个(19个类+ 1个背景)不同的值,而不是35个。

def encode_labels(mask):
    label_mask = np.zeros_like(mask)
    for k in mapping_20:
        label_mask[mask == k] = mapping_20[k]
    return label_mask

然后,您可以使用这些新的类数正常训练模型。