如何创建具有多个标签和掩码的自定义 Pytorch 数据集？答案

【问题标题】：How to Create a Custom Pytorch Dataset with Multiple Labels and Masks?如何创建具有多个标签和掩码的自定义 Pytorch 数据集？
【发布时间】：2021-07-16 00:28:46
【问题描述】：

我正在尝试创建一个自定义pytorch 数据集以插入DataLoader，该数据集由单通道图像(20000 x 1 x 28 x 28)、单通道掩码(20000 x 1 x 28 x 28) 和三个标签(20000 X 3) 组成。

按照文档，我想我会使用以下代码测试创建一个具有单通道图像和单通道掩码的数据集：

class CustomDataset(Dataset):

    def __init__(self, images, masks, transforms=None, train=True): 
        self.images = images
        self.masks = masks
        self.transforms = transforms

    def __getitem__(self, index):
        image = self.images.iloc[index, :]
        image = np.asarray(image).astype(np.float64).reshape(28, 28, 1)
        mask = self.masks.iloc[index, :]
        mask = np.asarray(mask).astype(np.float64).reshape(28, 28, 1)
        transformed_image = self.transforms(image)
      
        return transformed_image, mask

    def __len__(self):  

       return len(self.images)

使用该类，我从两个 pandas 数据帧形成数据集并插入DataLoader。

transform = transforms.Compose(
    [transforms.ToPILImage(),
     transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))
])

train_images = pd.read_csv('train.csv')
train_masks = pd.read_csv('masks.csv')

train_data = CustomDataset(train_images, train_masks, transform)
trainloader = DataLoader(train_data, batch_size=128, shuffle=True)

我希望trainloader 中的单个批次的形状为([128, 1, 28, 28], [128, 1, 28, 28])，对于左侧的图像和右侧的蒙版。

而不是单个批次trainloader的形状是([128, 1, 28, 28], [128])，这让我觉得掩码不知怎么变成了标签。

如何解决这个问题，除了掩码之外，如何添加实际标签？提前感谢您的帮助！

【问题讨论】：

标签： python deep-learning pytorch dataloader pytorch-dataloader

【解决方案1】：

也许您也需要在蒙版上应用变换（不包括标准化）。喜欢

class CustomDataset(Dataset):

def __init__(self, images, masks, transforms_image=None, transforms_mask=None, train=True): 
    self.images = images
    self.masks = masks
    self.transforms_image = transforms_image
    self.transforms_mask = transforms_mask
def __getitem__(self, index):
    image = self.images.iloc[index, :]
    image = np.asarray(image).astype(np.float64).reshape(28, 28, 1)
    mask = self.masks.iloc[index, :]
    mask = np.asarray(mask).astype(np.float64).reshape(28, 28, 1)
    transformed_image = self.transforms(image)
  
    return transformed_image, mask

def __len__(self):  

   return len(self.images)

和

   transform_image = transforms.Compose(
    [transforms.ToPILImage(),
     transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))
])

   transform_mask = transforms.Compose(
    [transforms.ToPILImage(),
     transforms.ToTensor()
])
train_images = pd.read_csv('train.csv')
train_masks = pd.read_csv('masks.csv')

train_data = CustomDataset(train_images, train_masks, transform_image, transform_mask)
trainloader = DataLoader(train_data, batch_size=128, shuffle=True)

【讨论】：