【发布时间】: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?如果你想让它更快,我猜你应该增加这个数字
标签: python pytorch gpu dataloader