如何将 cityscapes 数据集的 35 类转换为 19 类？

Question

以下是我的一小段代码。使用它，我可以在 cityscapes 数据集上训练名为 'lolnet' 的模型。但是数据集包含 35 classes/labels [0-34]。

imports ***

trainloader = torch.utils.data.DataLoader(
    datasets.Cityscapes('/media/farshid/DataStore/temp/cityscapes/', split='train', mode='fine',
                    target_type='semantic', target_transform =trans,
                    transform=input_transform ), batch_size = batch_size, num_workers = 2)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = lolNet()
criterion = CrossEntropyLoss2d()

net.to(device)
num_of_classes = 34

for epoch in range(int(0), 200000):

    lr = 0.0001

    for batch, data in enumerate(trainloader, 0):

        inputs, labels = data
        labels = labels.long()
        inputs, labels = inputs.to(device), labels.to(device)

        labels = labels.view([-1, ])

        optimizer = optim.Adam(net.parameters(), lr=lr)

        optimizer.zero_grad()
        outputs = net(inputs)

        outputs = outputs.view(-1, num_of_class)


        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        outputs = outputs.to('cpu')
        outputs = outputs.data.numpy()
        outputs = outputs.reshape([-1, num_of_class])

        mask = np.zeros([outputs.shape[0]])
        #
        for i in range(len(outputs)):
            mask[i] = np.argmax(outputs[i])

        mask = mask.reshape([-1, 1])

        IoU = jaccard_score(labels.to('cpu').data, mask, average='micro')

但我只想在 19 类上训练我的模型。这 19 类被发现 here . The labels to train for are stored as "ignoreInEval" = True. This pytorch Dataloader helper 对于这个数据集没有提供任何线索。

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

Answer 1

您下载模型和权重。

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

r = models.resnet50(pretrained=True)

注意原来resent有1000categories/classes。因此，当您下载预训练模型时，最后一个 fc 将用于 1000 类.

这是您拥有的 forward() 方法，上面的代码是您的模型。

您可以从原始 resnet50 模型中删除最后一个 fc 全连接层，并添加新的 fc 恰好有 19 类（19 个输出），然后您可以训练分类器仅适用于最后一层。其他层，除了最后一层应该被冻结。

所以您将只学习您需要的 19 类。

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

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

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

Answer 2

已经过了一段时间，但留下一个对其他人有用的答案：

首先创建一个到 19 classes + 背景的映射。背景与不那么重要的 classes 有关，带有忽略标志，如 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
    }

然后对于您加载的每个标签图像（每个像素包含 class 的灰色图像，其具有模式“{city}__{number}_{number}_gtFine_labelIds.png”）对于训练，运行下面的函数。

它将根据上面的映射转换每个像素，您的标签图像（蒙版）现在只有 20（19 classes + 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

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

如何将 cityscapes 数据集的 35 类转换为 19 类？

How to convert 35 classes of cityscapes dataset to 19 classes?

computer-vision

image-segmentation

pytorch

如何将 cityscapes 数据集的 35 类 转换为 19 类？

How to convert 35 classes of cityscapes dataset to 19 classes?

computer-vision

image-segmentation

pytorch

如何将 cityscapes 数据集的 35 类转换为 19 类？