KnowHow

PythonでGPUを使ったテスト計算のサンプル（PyTorchとTensorFlowの両方）の簡単なコード例を紹介します。

PyTorch版（GPUテスト）
python

import torch

# GPUが使えるか確認
print("GPU available:", torch.cuda.is_available())

# テンソルをGPUに移動して計算
x = torch.rand(10000, 10000).to('cuda')
y = torch.rand(10000, 10000).to('cuda')
z = x @ y  # 行列積

print("計算終了。z.shape =", z.shape)

torch.cuda.is_available() でCUDA GPUが認識されるか判定します。

GPUが使える場合、テンソル（x, y）を .to('cuda') でGPUに移動し、行列積（@演算子）で負荷のある計算を行います。

TensorFlow版（GPUテスト）

import tensorflow as tf

# GPUが使えるか確認
print("GPU available:", tf.config.list_physical_devices('GPU'))

# テンソルを作ってGPU上で計算
with tf.device('/GPU:0'):
    a = tf.random.uniform([10000, 10000])
    b = tf.random.uniform([10000, 10000])
    c = tf.matmul(a, b)  # 行列積

print("計算終了。c.shape =", c.shape)

tf.config.list_physical_devices('GPU') で利用可能なGPU一覧を表示できます。

with tf.device('/GPU:0'): で明示的にGPU上で計算させています.

どちらのサンプルも「大きな行列積」を計算することでGPUの利用状況や速度向上が確認できます。
必要なライブラリ（PyTorch, TensorFlow）はインストール済みであることを前提としています。

GPUテストに使うなら、このような二行列の積が最もシンプルかつ負荷をかけやすいです。出力を見て「GPU available: True」「c.shape = (10000, 10000)」等が出れば正常にGPUを利用できています。

上記のPyTorchおよびTensorFlowテストプログラムは、
「明示的にどのGPUを使うか指定しなければ、通常1個のGPU（PyTorch:デフォルトでcuda:0，TensorFlow:/GPU:0）だけを利用します」。

詳細
PyTorch

.to('cuda')はcuda:0（1番目のGPU）を利用します。

8枚GPUがあっても指定しない限り1枚のみ（自動では複数枚は使いません）。

TensorFlow

with tf.device('/GPU:0'):と書くことで、0番目のGPUだけ使います。

指定しない場合もデフォルトで1枚のみ割り当てられます。

複数GPUを使うには？
PyTorchではtorch.nn.DataParallelやtorch.nn.parallel.DistributedDataParallelなどを使い、

TensorFlowでもtf.distribute.MirroredStrategy等で複数GPU分散処理ができますが、

上記サンプルのような単純計算の場合は「1枚のみ」動作です。

このコードの計算時間は実行環境のGPU性能（種類・世代）、ドライバー・CUDA・TensorFlowバージョン、空きメモリ、その他の負荷によって大きく変わります。

おおよその目安
最新のNVIDIA A100やH100クラスのGPUなら、数秒～10秒以内で終わる可能性があります。

古いGPU（Tesla K80・V100・P100など）や物理的なマシン負荷が高い場合、10秒～1分程度かかる場合もあります。

CPUで同じ計算をすると数分～数十分以上かかることもあるので、GPU上での行列積は極めて高速です。

重要なポイント
本当に厳密な計算時間を知りたい場合は、timeライブラリで実時間計測が必要です。

例（TensorFlow）：

import tensorflow as tf
import time

with tf.device('/GPU:0'):
    a = tf.random.uniform([10000, 10000])
    b = tf.random.uniform([10000, 10000])

start = time.time()
c = tf.matmul(a, b)
tf.keras.backend.clear_session()  # メモリ解放
print("経過時間(sec):", time.time() - start)

まとめ
一般的なGPU（V100, A100, RTX3090等）なら「数秒～数十秒」程度で終わるケースが多いです。
厳密な時間は環境に依存するので、コード計測で実測するのがベストです。
「数分かかる」ことはGPUであればほぼなく、CPUだと数分～数十分が目安です。

ユーザー指定で使用GPU数を変更する一般的な方法は、
「環境変数を設定」「APIパラメータで明示する」「スクリプト引数で受けて制御する」の3つが多いです。

1. CUDA_VISIBLE_DEVICES環境変数（PyTorch/TensorFlow共通）

# 例: 2枚だけ使いたい場合
export CUDA_VISIBLE_DEVICES=0,1
python your_script.py

0,1 → 1番目と2番目のGPUのみ認識され、プログラムからは2枚だけ見える。

1. TensorFlow（API側で枚数指定や分散戦略）
   python

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
# 使いたいGPU数を指定
num_gpus_to_use = 2
tf.config.set_visible_devices(gpus[:num_gpus_to_use], 'GPU')

上記で使う枚数や組み合わせをプログラム中で柔軟に指定できます。

1. PyTorch（DataParallelやDistributedDataParallel）

import torch
import torch.nn as nn

device_ids = [0, 1, 2]  # 使いたいGPUのIDリスト

model = nn.DataParallel(YourModelClass(), device_ids=device_ids)
model = model.to('cuda:0')

device_idsをコマンドライン引数や設定ファイルから可変にして制御可能。

まとめ
どの方法でも「見えるGPU枚数」を環境変数やAPIで制御可能。

TensorFlowならtf.config.set_visible_devices、PyTorchならdevice_ids=[...]が主力。

コマンドライン引数や設定ファイルからユーザーが希望した枚数を受け取り、その値を反映する実装パターンが推奨です。

「環境変数による一括制御」でも「プログラム内で欲しい枚数だけ利用」でもOK。どちらも現場でよく使われています。

このコードを実行すると、4枚のGPUを使った計算が走り、計算時間と結果の形状が出力されます。
環境によってはdevice_idsの枚数や番号を調整してください。4枚のGPUを使ってPyTorchでテスト計算する簡単なサンプルコードを示します。
ここではnn.DataParallelを使い、4つのGPUを指定して行列積をGPU並列で計算します。

import torch
import torch.nn as nn
import time

# 4つのGPUを使う設定
device_ids = [0, 1, 2, 3]

# サンプルモデル（単純な行列積演算用）
class MatrixMulModel(nn.Module):
    def forward(self, x, y):
        return torch.matmul(x, y)

# モデル作成＆DataParallelでマルチGPU対応に
model = MatrixMulModel()
model = nn.DataParallel(model, device_ids=device_ids)
model = model.cuda(device_ids[0])  # メインGPU0にモデルを置く

# 大きい行列を用意し、GPUへ移動
x = torch.rand(10000, 10000).cuda(device_ids[0])
y = torch.rand(10000, 10000).cuda(device_ids[0])

# 時間計測開始
torch.cuda.synchronize()  # GPU同期（計測精度向上）
start_time = time.time()

# 計算（モデルを使って行列積）
with torch.no_grad():
    z = model(x, y)

torch.cuda.synchronize()  # 終了同期
end_time = time.time()

print(f"計算時間: {end_time - start_time:.4f}秒")
print(f"結果の形状: {z.shape}")

ポイント
device_idsで使いたいGPUのIDをリスト指定。今回は4つのGPU (0,1,2,3)を指定

nn.DataParallelがGPU間でモデルと計算を並列処理してくれる

テンソルは計算開始前に最初のGPU（ID0）に移動させる

計算前後に torch.cuda.synchronize() で同期して正確に計測

単純な行列積だが4枚のGPUを使ってパフォーマンス向上が期待できる

このコードを実行すると、4枚のGPUを使った計算が走り、計算時間と結果の形状が出力されます。
環境によってはdevice_idsの枚数や番号を調整してください。

以下は、PyTorchの4GPU利用テスト計算スクリプトをSLURMで実行するためのシンプルなバッチスクリプト例です。

#!/bin/bash
#SBATCH --job-name=pytorch_gpu_test      # ジョブ名
#SBATCH --partition=gpu                   # GPU用パーティション（環境により変更）
#SBATCH --nodes=1                        # ノード数
#SBATCH --ntasks=1                       # タスク数
#SBATCH --cpus-per-task=4                # CPUコア数
#SBATCH --gres=gpu:4                     # GPUを4枚要求
#SBATCH --time=01:00:00                  # 実行時間制限（1時間）
#SBATCH --output=pytorch_gpu_test_%j.out # 標準出力ログ
#SBATCH --error=pytorch_gpu_test_%j.err  # 標準エラーログ
#SBATCH --mail-type=END,FAIL             # ジョブ終了時メール
#SBATCH --mail-user=your_email@example.com  # メールアドレス

# モジュールロード（環境により）
module purge
module load cuda
module load pytorch/2.0.0  # PyTorchモジュールがあれば指定

# 作業ディレクトリ移動（必要に応じて）
cd $SLURM_SUBMIT_DIR

# Pythonスクリプト実行
srun python your_pytorch_gpu_test.py

補足
--gres=gpu:4 で4枚のGPUを要求しています

srun python your_pytorch_gpu_test.py の your_pytorch_gpu_test.py に先ほどの4GPU用PyTorchテストコードを置いてください

環境により module load の行は調整してください（CUDAやPyTorchが使えるように）

メール通知やログ出力パスも必要に応じて編集してください

このスクリプトを例として sbatch run_gpu_test.slurm のように投入すれば、4GPU環境でテスト計算が動作します。

--cpus-per-task の指定は必須ではありませんが、以下のような理由で指定することが多いです。

意味と役割
--cpus-per-task は「1タスク（プロセス）あたりに割り当てるCPUコア数」を指定します。

CPUマルチスレッド処理や並列化を想定したプログラムで、1つのタスクが複数CPUコアを使う場合に設定します。

指定しないと、Slurmの設定やパーティションのデフォルトのコア数が割り当てられることがあります（少ない場合もあり）。

GPUジョブでの実際
GPUを使った処理は多くの場合CPUは補助的なので、CPUコアは最低限（例えば4や8）で十分なことが多いです。

一方で、データロードや前処理でマルチスレッドを使うなら、--cpus-per-task は適宜増やします。

指定しないと「CPUコアが1コアだけ」割り当てられて処理が遅くなったり、ジョブの待ち時間が長くなることもあります。

まとめ
プログラムが1タスク中に複数CPUコアを使うなら指定推奨

単純にGPUだけ使うだけなら最低値で問題なし（ただし4くらいは割り当てたほうが安全）

指定しないと管理側のデフォルトコア数になる（要注意）

基本的には負荷に応じて適宜調整してください。
たとえば、4枚GPU使うなら少なくとも--cpus-per-task=4くらいはあったほうが無難です。

srun --gpus=4 run.sh のようにGPU数だけ指定した場合、SLURMはCPUコア数はデフォルトで1つ割り当てるケースが多いです。

そのため、CPUを複数コア使いたい場合は、明示的に --cpus-per-task= などで指定するのが良いです。

理由とポイント
GPU数を指定してもCPUコア数は自動で増えないので、CPUリソースが不足する可能性がある

データロードや前処理・マルチスレッド処理を行う場合はCPUコア数を増やさないと全体の速度が落ちる

ジョブスケジューラはユーザーのリクエストに応じてCPUを割り当てるため、過不足なく明示指定するのがベストプラクティス

例

#SBATCH --gres=gpu:4
#SBATCH --cpus-per-task=8

などとしてCPU/メモリリソースを増やしておくことで、GPUとCPUのバランスをうまく取れます。

総じて「GPU数は指定したがCPUコアは指定していなかった」ケースでCPUが1コア割当てられ、ジョブ性能が出せなかった現象はよくありますので、
GPUを複数使うならCPUも適切に割り当ててくださいが運用上の鉄則です。

SLURMでGPUを複数割り当てた場合、ジョブ内のプログラムは環境変数を通じて割り当てられたGPU情報を受け取り、それを利用してどのGPUを使うか認識します。

主な認識方法
SLURMは、割り当てたGPUを環境変数 CUDA_VISIBLE_DEVICES に書き込みます。

これにより、プログラム内のCUDA / PyTorch / TensorFlowなどのGPU管理層は、CUDA_VISIBLE_DEVICESが示すGPUだけを見るようになります。

例: 環境変数の見え方
SLURMジョブ開始後のNode内で、

echo $CUDA_VISIBLE_DEVICES

とすると、

0,1,2,3

のように割り当てGPUのIDがカンマ区切りで表示されます。（例えばGPU ID 0〜3の4枚が割り当てられた場合）

プログラム内でのGPU認識
PyTorch や TensorFlow は内部で CUDA_VISIBLE_DEVICES を読み込んで見えるGPUのリストを決定します。

例えば PyTorchなら torch.cuda.device_count() はこの環境変数に基づいたGPU数を返します。

TensorFlowでも同様に tf.config.list_physical_devices('GPU') は割り当てられたGPUリストになります。

そのためスクリプトに特別な記述は不要

とりあえず、SLURMでGPUが参照できているかを確かめるスクリプトサンプル

#!/bin/bash
#SBATCH --job-name=pytorch_gpu_test
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:4
#SBATCH --time=00:10:00
#SBATCH --output=pytorch_gpu_test_%j.out
#SBATCH --error=pytorch_gpu_test_%j.err

# module load cuda

echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"

nvidia-smi

sleep 120