GPUの使用方法

[更新: 2024年06月27日]

※表記やコードに誤りがあれば適宜修正お願いします

GPUアーキテクチャ

高火力 DOK では、NVIDIA社製のGPUを提供しています。 そのため、同社の開発・提供するCUDAを使うことでGPUを使用できます。

コンテナー内にCUDAがインストールされている必要がありますので、 CUDAインストール済みのコンテナーイメージ を元に構築するか、 ご自分でCUDAをインストールしたコンテナーイメージをご使用ください。

ソースコード例

PyTorch

PyTorchでGPUを使用するコードの例です。

# PyTorchでGPUを使用する例
import torch

t = torch.tensor([0, 1, 2], device=torch.device('cuda'))

プログラムを本サービス以外（CUDAを使えないマシン）でも動かす可能性がある場合は、以下のように記述するとCUDAを使えない場合にCPUを使います。

# CUDAを使えない場合はCPUを使う例
import torch

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
t = torch.tensor([0, 1, 2], device=device)