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

Question

在训练循环中，我将一批数据加载到 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 而无需传输步骤？

Answer 1

您可以将所有数据加载到张量中，而不是将其移动到 GPU 内存中。（假设您有足够的内存）当您需要它时，使用已经在 GPU 内存中的张量中的数据。希望对你有帮助。

Answer 2

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

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

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

Load data into GPU directly using PyTorch

gpu

python-3.x

deep-learning

pytorch

dataloader