我正在尝试创建一个自定义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])
,这让我觉得面具已经以某种方式变成了标签。
我该如何解决这个问题,除了掩码之外,我如何添加实际标签?在此先感谢您的帮助!