将 pytorch 数据加载器加载到 GPU 中

Posted 2023-02-16

【中文标题】将 pytorch 数据加载器加载到 GPU 中

【英文标题】：load pytorch dataloader into GPU

【发布时间】：2021-03-27 08:22:48

【问题描述】：

有没有办法将 pytorch DataLoader (torch.utils.data.Dataloader) 完全加载到我的 GPU 中？

现在，我将每个批次分别加载到我的 GPU 中。

CTX = torch.device('cuda')

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=0,
)

net = Net().to(CTX)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=LEARNING_RATE)

for epoch in range(EPOCHS):
    for inputs, labels in test_loader:
        inputs = inputs.to(CTX)        # this is where the data is loaded into GPU
        labels = labels.to(CTX)        

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    print(f'training accuracy: net.validate(train_loader, device=CTX)/len(train_dataset)')
    print(f'validation accuracy: net.validate(test_loader, device=CTX)/len(test_dataset)')

Net.validate() 函数由下式给出

def validate(self, val_loader, device=torch.device('cpu')):
    correct = 0
    for inputs, labels in val_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = torch.argmax(self(inputs), dim=1)
        correct += int(torch.sum(outputs==labels))
    return correct

我想通过将整个数据集 trainloader 加载到我的 GPU 中来提高速度，而不是单独加载每个批次。所以，我想做类似的事情

train_loader.to(CTX)

这有等效的功能吗？因为torch.utils.data.DataLoader没有.to()这个属性。

我使用安装了 CUDA Toolkit 10.2 的 NVIDIA GeForce RTX 2060。

【问题讨论】：

为什么将num_workers 设置为0？如果你想让它更快，我猜你应该增加这个数字

【参考方案1】：

你可以提前把数据集的数据放好

train_dataset.train_data.to(CTX)  #train_dataset.train_data is a Tensor(input data)
train_dataset.train_labels.to(CTX)

例如minst

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
batch_size = 64
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_data = datasets.MNIST(
    root='./dataset/minst/',
    train=True,
    download=False,
    transform=transform
)
train_loader = DataLoader(
    dataset=train_data,
    shuffle=True,
    batch_size=batch_size
)
train_data.train_data.to(torch.device("cuda:0"))  # put data into GPU entirely
train_data.train_labels.to(torch.device("cuda:0"))

我通过使用调试器得到了这个解决方案...

【参考方案2】：

在本教程 (https://pytorch.org/tutorials/beginner/nn_tutorial.html) 的“包装数据加载器”部分中，数据被完全加载到 GPU 中。包装器数据加载器代码如下：

def preprocess(x, y):
    return x.view(-1, 1, 28, 28).to(dev), y.to(dev)

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)