DataLoader转GPU
时间: 2024-02-05 12:04:20 浏览: 26
要将DataLoader转移到GPU上,需要将数据和标签转换为torch张量,然后使用.to()方法将其移动到GPU上。以下是一个示例:
```python
import torch
from torch.utils.data import DataLoader
# 定义数据集
class MyDataset(torch.utils.data.Dataset):
def __init__(self, data, labels):
self.data = data
self.labels = labels
def __getitem__(self, index):
x = self.data[index]
y = self.labels[index]
return x, y
def __len__(self):
return len(self.data)
# 创建数据集
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
labels = [0, 1, 0]
dataset = MyDataset(data, labels)
# 创建DataLoader
batch_size = 2
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 将数据转换为GPU张量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for batch_idx, (data, labels) in enumerate(dataloader):
data, labels = data.to(device), labels.to(device)
# 在此处进行模型训练
```
在上面的示例中,我们首先定义了一个简单的数据集类,然后使用它创建了一个DataLoader。接下来,我们使用torch.device()方法检查GPU是否可用,并在需要时将数据和标签转换为GPU张量。最后,我们可以使用GPU上的数据和标签进行模型训练。
相关推荐
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)