Pytorch 张量使用所有 RAM

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