[PyTorch] GPUの動作確認セット

はじめに

私は普段、PyTorchの環境構築をしたあとに確認用のコードを探して導入確認しています。ですが、いつも複数のサイトを探す時間がもったいないのと結局私好みの出力が欲しかったため、この場所に動作確認セットおよび備忘録として残しておきます。

動作確認セット

import torch

print(torch.__version__)
print(f"cuda, {torch.cuda.is_available()}")
print(f"compute_{''.join(map(str,(torch.cuda.get_device_capability())))}")
device_num:int = torch.cuda.device_count()
print(f"find gpu devices, {device_num}")
for idx in range(device_num):
    print(f"cuda:{idx}, {torch.cuda.get_device_name(idx)}")

print("end")

出力例

1.8.0+cu101
cuda, True
compute_61
find gpu devices, 1
cuda:0, GeForce GTX 1080
end

最後に

これで自分があちこち検索しに回らなくて済むかも・・・