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

Question

有没有办法将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。

Answer 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"))

我使用调试器得到了这个解决方案。。。

Answer 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)