データをGPUに送りたい

今日は何をしたい？

CPU にある tensor を GPU へ送って、device が変わる様子を見ます。計算そのものより、まずは「置き場所」を読めることが目標です。

確認コード

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)

print("before:", x.device)
x = x.to(device)
print("after :", x.device)

このコードを chapter05.py として保存して、python chapter05.py で実行します。

GPU 環境なら、こう表示されます。

GPU がない環境では after も cpu のままになります。それはエラーではなく、CPU で実行されている正常な状態です。

before と after を見比べるだけで、データがどこにあるかを追えるようになります。

どこがROCm？

ROCm 版 PyTorch でも、GPU を指す場所で "cuda" という名前が出ます。名前に引っ張られず、「GPU 側へ送るための入口」と読むのがコツです。

出力の cuda:0 は、ROCm 経由で接続されている AMD GPU の 0 番を指しています。第1章で torch.cuda.is_available() が True になっていた GPU が、ここで初めて実際にデータを受け取ります。

ここで出てきたPython

• "A" if 条件 else "B" — 三項演算子で分岐。条件が True なら A、そうでなければ B を返します。この章では「GPU があれば "cuda"、なければ "cpu"」を1行で書くために使っています。以降の章でも同じ書き方が繰り返し登場します。
• .to(device) — データの置き場所を変えるメソッド

よくあるつまずき

• tensor だけ GPU に送っても、モデルが CPU のままだと次の章で混ざって止まることがあります。
• device が cpu のままでもエラーではありません。まずは環境の状態を知ることが先です。
• "cuda" という名前を見て混乱しやすいですが、ROCm でもここはよく出てきます。

device が混ざるとどうなるか、見てみましょう。

# device 不一致の例
import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Linear(2, 1)          # CPU のまま
x = torch.tensor([[1.0, 2.0]], device=device)  # GPU

try:
    model(x)
except RuntimeError as e:
    print(e)

1分演習

コードの最後に print(x) を足して、値そのものは同じでも置き場所だけが変わることを見てみましょう。