pytorch使用GPU
时间: 2024-01-27 20:15:20 浏览: 47
在PyTorch中使用GPU可以显著加速模型的训练和推理过程。下面是两种使用GPU的方法:
1. 使用`.to()`方法将模型和数据移动到GPU上:
```python
import torch
# 检查是否有可用的GPU
if torch.cuda.is_available():
device = torch.device("cuda") # 使用默认的GPU设备
else:
device = torch.device("cpu") # 如果没有GPU,则使用CPU
# 定义模型
model = MyModel()
model.to(device) # 将模型移动到GPU上
# 定义数据
inputs = torch.randn(10, 3).to(device) # 将输入数据移动到GPU上
labels = torch.randn(10, 1).to(device)
# 在GPU上进行训练
for epoch in range(num_epochs):
# 前向传播和反向传播
outputs = model(inputs)
loss = loss_function(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
```
2. 使用`torch.nn.DataParallel`在多个GPU上并行训练模型:
```python
import torch
import torch.nn as nn
from torch.nn import DataParallel
# 定义模型
model = MyModel()
# 检查是否有可用的多个GPU
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model) # 在多个GPU上并行训练模型
model.to(device) # 将模型移动到GPU上
# 定义数据
inputs = torch.randn(10, 3).to(device) # 将输入数据移动到GPU上
labels = torch.randn(10, 1).to(device)
# 在GPU上进行训练
for epoch in range(num_epochs):
# 前向传播和反向传播
outputs = model(inputs)
loss = loss_function(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
```