要在PyTorch中實現GPU加速,首先確保安裝了支持GPU的PyTorch版本。然后可以通過以下步驟在GPU上運行PyTorch代碼:
import torch
if torch.cuda.is_available():
device = torch.device("cuda")
print("GPU is available")
else:
device = torch.device("cpu")
print("GPU is not available, using CPU instead")
model = YourModel().to(device)
input_tensor = torch.randn(1, 3, 224, 224).to(device)
output = model(input_tensor)
model.parameters()
和optimizer
的step()
方法:optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
optimizer.zero_grad()
output = model(input_tensor)
loss = loss_function(output, target)
loss.backward()
optimizer.step()
通過以上步驟,就可以在PyTorch中實現GPU加速,并利用GPU的并行計算能力加快模型訓練和推理的速度。