如何将dataloader的数据移动到cuda上
时间: 2024-03-08 11:51:08 浏览: 23
将 `DataLoader` 中的数据移动到 GPU 上,可以使用 `to()` 方法。具体来说,可以在训练循环中使用 `to()` 方法将输入数据和标签数据移动到 GPU 上。
示例代码如下:
```
import torch
from torch.utils.data import DataLoader
# 定义训练数据集
train_dataset = ...
# 定义 DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# 定义模型并移动到 GPU 上
model = MyModel()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# 定义损失函数和优化器
criterion = ...
optimizer = ...
# 训练循环
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(train_loader):
# 将输入数据和标签数据移动到 GPU 上
inputs = inputs.to(device)
labels = labels.to(device)
# 前向传播、计算损失、反向传播、更新参数
...
```
在这个示例代码中,我们首先定义了一个 `train_dataset` 和一个 `train_loader`,然后定义了一个模型 `model` 并将其移动到 GPU 上。在训练循环中,我们使用 `to()` 方法将每个 batch 的输入数据和标签数据移动到 GPU 上,然后进行前向传播、计算损失、反向传播和更新参数的操作。