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

Posted 2023-02-16

【中文标题】使用 PyTorch 直接将数据加载到 GPU 中

【英文标题】：Load data into GPU directly using PyTorch

【发布时间】：2020-09-18 12:52:43

【问题描述】：

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

【问题讨论】：

你凭什么说它很耗时？我认为这是标准做法。

【参考方案1】：

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

【讨论】：

@PeterJulian 准备好包含数据的张量后，您可以使用 your_data = your_data.to(device) 将其移动到 GPU 您可以找到一些示例和详细信息 here。

谢谢，我知道你可以用它把张量加载到设备上。我的意思是，是否有任何命令可以将整个数据集加载到 GPU 中，这样您就不必在每批中调用（设备）。不确定这有多昂贵，但它始终是从 CPU 到 GPU 的操作，这在计算中可能很明显，尤其是。在较小的网络/数据集上。

【参考方案2】：

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

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