PyTorchでGPUアクセラレーションを実現する方法は何ですか？

1年 ago
海斗, 葵
1 minute

PyTorchをGPUで実行するためには、まずGPUをサポートしているPyTorchバージョンをインストールしてください。そして、以下の手順に従ってGPU上でPyTorchコードを実行できます。

  1. GPUデバイスが利用可能かどうかをチェックしてください。
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")
  1. GPUデバイスにモデルとテンソルを移動します。
model = YourModel().to(device)
input_tensor = torch.randn(1, 3, 224, 224).to(device)
  1. GPUを使用してモデルのトレーニングと推論を行う：
output = model(input_tensor)
  1. モデルのパラメータ()
  2. 最適化プログラム
  3. ステップ()
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の並列計算能力を活用してモデルの学習と推論の速度を向上させることができます。

bannerAds