使用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内存中的张量中的一个。希望对您有所帮助。