pytorch自定义数据集:DataLoader返回张量列表,而不是列表的张量

时间:2020-06-05 05:44:50

标签: python pytorch

import torch

class Custom_Dataset(torch.utils.data.dataset.Dataset):
    def __init__(self, _dataset):
        self.dataset = _dataset

    def __getitem__(self, index):
        example, target = self.dataset[index]
        return example, target

    def __len__(self):
        return len(self.dataset)

train_data = [([1, 3, 5], 0),
              ([2, 4, 6], 1)]
train_loader = torch.utils.data.DataLoader(dataset=Custom_Dataset(train_data),
                                           batch_size=1,
                                           shuffle=False)

for inputs, targets in train_loader:
    print(inputs)
    print(targets)

我将训练数据定义为[([1, 3, 5], 0), ([2, 4, 6], 1)]:输入([1, 3, 5])配对目标(0)。

但是当我从数据加载器中获取数据时,它变为:

[tensor([1]), tensor([3]), tensor([5])]
tensor([0])

我该如何获取:

tensor([[1],
        [3],
        [5]])
tensor([0])

我知道torch.stack可以解决问题,但是我可以在自定义数据集类中将其转换吗?

1 个答案:

答案 0 :(得分:1)

一种获得所需输入的解决方案是使用numpy。下面,我仅更改了示例中的两行以使其起作用。

import torch
import numpy as np

class Custom_Dataset(torch.utils.data.dataset.Dataset):
    def __init__(self, _dataset):
        self.dataset = _dataset

    def __getitem__(self, index):
        example, target = self.dataset[index]
        return np.array(example), target

    def __len__(self):
        return len(self.dataset)

train_data = [([1, 3, 5], 0),
              ([2, 4, 6], 1)]
train_loader = torch.utils.data.DataLoader(dataset=Custom_Dataset(train_data),
                                           batch_size=1,
                                           shuffle=False)

for inputs, targets in train_loader:
    print(inputs)
    print(targets)

此代码的输出为

tensor([[1, 3, 5]])
tensor([0]) 

tensor([[2, 4, 6]])
tensor([1])

但是,当然,我假设拥有行向量或列向量对您没有任何影响。否则,您可能需要检查this answer about transposing 1D vectors

希望这会有所帮助。