MENU

Pytorch – 計算を行うデバイスを指定する方法について

ディープラーニング
Pytorch ディープラーニング
  1. ホーム
  2. 機械学習
  3. ディープラーニング
  4. Pytorch – 計算を行うデバイスを指定する方法について
目次
  1. 概要
  2. torch.device
  3. デバイスの作成方法
  4. テンソルのデバイスの指定方法
  5. デバイスを返す関数を作成する
  6. デフォルトのデバイスを設定する

概要

Pytorch で計算を指定したデバイス (CPU または GPU) で行う方法について解説します。

torch.device

torch.device は、デバイスを表すクラスです。 torch.Tensor クラスは、計算を実行するデバイスを設定として持っており、計算はそのデバイス上で行われます。テンソルのデバイスは Tensor.device で取得できます。

In [1]: 
import torch

a = torch.tensor([1, 2, 3], device=torch.device("cpu"))
b = torch.tensor([1, 2, 3], device=torch.device("cpu"))

c = a * b  # a, b の device はどちらも CPU なので、計算は CPU 上で行われる。
print(f"a.device: {a.device}")
print(f"b.device: {b.device}")
print(f"c.device: {c.device}")
Python
a.device: cpu
b.device: cpu
c.device: cpu

計算を実行する際に、使用するテンソルのデバイスはすべて一致している必要があります。テンソルによってデバイスが異なる場合は、エラーとなります。

In [2]: 
import torch

a = torch.tensor([1, 2, 3], device=torch.device("cpu"))
b = torch.tensor([1, 2, 3], device=torch.device("cuda"))

c = a * b  # エラー
Python

デバイスの作成方法

  • CPU
    • torch.device("cpu")
  • GPU
    • torch.device("cuda")
    • 複数枚の GPU がある場合、使用する GPU の ID を torch.device("cuda:<GPU ID>") または torch.device("cuda", index=<GPU ID>) と指定する。
    • GPU ID は nvidia-smi コマンドで確認できる。1枚目は ID=0、2枚目は ID=1、… となっている。
In [3]: 
import torch

# CPU
a = torch.tensor([1, 2, 3], device=torch.device("cpu"))

# GPU
a = torch.tensor([1, 2, 3], device=torch.device("cuda"))

# ID=0 の GPU (GPU が複数枚ある場合)
a = torch.tensor([1, 2, 3], device=torch.device("cuda:0"))
# または
a = torch.tensor([1, 2, 3], device=torch.device("cuda", index=0))
Python

デバイスを引数に指定する場合に、torch.device オブジェクトの代わりにデバイスを表す文字列を直接指定することもできます。

In [4]: 
import torch

# torch.device オブジェクトを指定する場合
a = torch.tensor([1, 2, 3], device=torch.device("cuda:0"))

# 文字列を指定する場合
a = torch.tensor([1, 2, 3], device="cuda:0")
Python

テンソルのデバイスの指定方法

計算を実行するデバイスの指定は、torch.tensor() でテンソルを作成する際に指定する方法と、すでにあるテンソルのデバイスを後から Tensor.to() で変更する方法があります。

  1. torch.tensor() でテンソルを作成する際に指定する。
  2. Tensor.to() で後から変更する。
In [5]: 
import torch

# torch.tensor() で作成時に device 引数に指定する。
a = torch.tensor([1, 2, 3], device="cuda")
print(a.device)  # cpu

# Tensor.to() で後から変更する。
a = torch.Tensor([1, 2, 3])
a = a.to("cuda")
print(a.device)  # cuda:0
Python
cuda:0
cuda:0

デバイスを返す関数を作成する

以下のような torch.device オブジェクトを返す関数を定義しておくと便利かもしれません。

  • CUDA が利用可能かどうか torch.cuda.is_available() でチェックする。
  • GPU を利用する場合は、gpu_id で使用する GPU ID を指定する。
  • CPU を利用する場合はなにも指定しない。
In [6]: 
def get_device(gpu_id=-1):
    if gpu_id >= 0 and torch.cuda.is_available():
        return torch.device("cuda", gpu_id)
    else:
        return torch.device("cpu")


device = get_device()
print(device)  # cpu

device = get_device(gpu_id=0)
print(device)  # cuda:0
Python
cpu
cuda:0

デフォルトのデバイスを設定する

テンソルを作成する際にデバイスを指定しない場合、set_default_tensor_type() で設定されているデバイスが設定されます。 デフォルトでは、CPU になっていますが、以下のように GPU で作成されるように変更できます。

In [7]: 
torch.set_default_tensor_type(torch.cuda.FloatTensor)
a = torch.tensor([1, 2, 3])
print(a.device)
Python
cuda:0
ディープラーニング
Pytorch ディープラーニング

コメント

コメント一覧 （0件）

  • 本サイトとても分かりやすくまとめていただき大変助かっております。
    現在、下記３記事を参考にさせていただき、大量の画像から物体検出をGPUで実施できるように取り組んでおります。
    ①YOLOv3 – 学習済みモデルで画像から人や車を検出する方法
    ②Pytorch – GPU と対応するドライバ、CUDA、CuDNN のバージョン
    ③Pytorch – 計算を行うデバイスを指定する方法について

    教えていただきたいのですが、②に従い、CUDA、CuDNN、Pytorchをセットアップすれば、①のコードについて自動でGPUにて実施されますでしょうか？
    ③で触れられているようにGPUにて演算させるPythonコード等を追記する必要はないでしょうか？

    理解不足の中での質問となり恐縮ですが、
    回答いただければ幸いです。

    返信

    • コメントありがとうございます。

      「Pytorch – GPU と対応するドライバ、CUDA、CuDNN のバージョン」ですが、内容が1年前のものとなっていたので更新しました。
      現在は、「CUDA 11.7」、「CuDNN8」、「最新の Nvidia ドライバ」がインストールされている環境で `pip install torch` でインストールされた最新の Pytorch が動くと思われます。(Pytorch のバージョンが古い場合は更新してください)

      「YOLOv3 – 学習済みモデルで画像から人や車を検出する方法」のコードでは、GPU が利用可能な環境では、自動で GPU が選択されるようになっているため、GPU 対応のためにコードを修正する必要はありません。
      お使いの環境で GPU が利用できる状態になっているかは以下のコードで確認できます。

      “`
      import torch
      print(torch.cuda.is_available())
      “`

      返信

  • ご回答いただいており感謝いたします。

    おかげさまで無事GPUでの処理ができるようになりました。

    改めてありがとうございました！

    返信

コメントする

目次
  1. 概要
  2. torch.device
  3. デバイスの作成方法
  4. テンソルのデバイスの指定方法
  5. デバイスを返す関数を作成する
  6. デフォルトのデバイスを設定する