使用所有 RAM 的 Pytorch 张量

Question

我有一个张量列表，这对我的 RAM 来说太重了。我想将它们保存在文件系统中并在需要时加载它们

torch.save(single_tensor, 'tensor_<idx>.pt')

如果我想在训练时使用批处理，有没有一种在需要时自动加载张量的方法？我正在考虑使用 TensorDataset 和 DataLoader，但现在我在列表中没有张量，但在文件系统中，我应该如何构建它们？

Answer 1

先将张量一个一个保存到文件torch.save()

torch.save(tensor, 'path/to/file.pt')

然后这个 Dataset class 允许仅在真正需要时加载张量：

class EmbedDataset(torch.utils.data.Dataset):
    def __init__(self, first_embed_path, second_embed_path, labels):
        self.first_embed_path = first_embed_path 
        self.second_embed_path = second_embed_path 
        self.labels = labels

        

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

    def __getitem__(self, i):

        label = self.labels[i]

        embed = torch.load(os.path.join(self.first_embed_path, str(i) + '.pt'))

        pos = torch.load(os.path.join(self.second_embed_path, str(i) + '.pt'))

        tensor = torch.cat((embed, pos))

        return tensor, label

这里张量用数字命名，例如1.pt或1816.pt

使用所有 RAM 的 Pytorch 张量

Pytorch Tensors using all RAM

python

pytorch