使用 PyTorch 直接将数据加载到 GPU 中

在训练循环中，我将一批数据加载到 CPU 中，然后将其传输到 GPU：

import torch.utils as utils

train_loader = utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)

for inputs, labels in train_loader:
    inputs, labels = inputs.to(device), labels.to(device)

这种加载数据的方式非常耗时。有什么方法可以直接将数据加载到 GPU 中而不需要传输步骤吗？

@PeterJulian 首先感谢您的回复。据我所知，没有单行命令可以将整个数据集加载到 GPU。实际上，在我的回复中，我的意思是在中使用 .to(device)__init__数据加载器的。我之前分享的链接中有一些示例。另外，我在下面留下了一个示例数据加载器代码。希望链接中的示例和下面的代码都有帮助。

class SampleDataset(Dataset):
    def __init__(self, device='cuda'):
        super(SampleDataset, self).__init__()
        self.data = torch.ones(1000)
        self.data = self.data.to(device)
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        element = self.data[i]
        return element