Pytorch：在图像预测中为图像添加信息

Question

我想向我当前的数据集添加信息。目前，我在文件夹中有六帧序列。 DataLoader 读取所有 6 个并使用前 3 个来预测最后一个 1/2/3（取决于我告诉他多少）。这是 DataLoader 的函数。

class TrainFeeder(Dataset):
    def init(self, data_set):
        super(TrainFeeder, self).init()
        self.input_data = data_set
        #print(torch.cuda.current_device())
        if torch.cuda.current_device() ==0:
            print('There are total %d sequences in trainset' % len(self.input_data))

    def getitem(self, index):
        path = self.input_data[index]
        imgs_path = sorted(glob.glob(path + '/*.png'))
        imgs = []
        for img_path in imgs_path:
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (256,448))
            img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC) #has been 0.5 for official data, new is fx = 2.63 and fy = 2.84
            img_tensor = ToTensor()(img).float()
            imgs.append(img_tensor)
        imgs = torch.stack(imgs, dim=0)
        return imgs

    def len(self):
        return len(self.input_data)

现在我想为这些图像添加一个值。它是一个布尔值，我存储在同一个文件夹中的 .json 列表中，就像六帧序列一样。但我不知道如何将 .json 中列表的值添加到张量中。我应该使用哪个维度？如果我改变输入的形状，系统会正常工作吗？

Answer 1

getitem 函数可以返回任何内容，因此您可以返回元组而不仅仅是图像：

def __getitem__(self, index):
    path = ...
    # load your 6 images
    imgs = torch.stack( ... )
    # load your boolean metadata
    metadata = load_json_data( ... )
    # return them both
    return (imgs, metadata)

您需要在返回之前将 metadata 设为张量，否则我希望 pytorch 会抱怨无法整理（即堆叠）它们以进行批量处理

“系统能否正常工作”是一个只有您可以回答的问题，因为您没有提供 ML 模型的代码。我敢打赌：“不，但它不需要进行重大更改即可工作”。很可能您目前有一个类似

的循环

for imgs in dataloader:
    # do some training 
    output = model(imgs)
    ...

你必须让它像

for imgs, metadata in dataloader:
     # do some training
     output = model(imgs)
     ...