CUDA・cuDNN・PyTorch・TensorFlowなどのバージョンを一括で表示してくれるPythonプログラム作ってみた【Python】

2024年10月2日Python,データ操作

生成AI関連の開発をしていると、バージョン確認が多く発生して面倒ですよね。

コマンドプロンプトでイチイチ叩いていくのも良いですが、ポチッと動かすだけで表示してくれる方が楽で助かるので、今回はそんなコードを作ってみました。リンクさえ貼っていただければ、二次配布もOKです。

1. バージョン確認プログラム

バージョン確認プログラム

コード

簡単ですが、以下がソースコードになります。

import torch
import subprocess
import importlib

def get_cuda_version():
    try:
        # コマンドラインでnvccのバージョンを取得
        output = subprocess.check_output(['nvcc', '--version']).decode('utf-8')
        for line in output.split('\n'):
            if 'release' in line:
                return line.split('release')[-1].strip()
    except Exception as e:
        return f"Error: {str(e)}"

def get_cudnn_version():
    try:
        # PyTorchがインストールされている場合、cudnnのバージョンを取得
        return torch.backends.cudnn.version()
    except Exception as e:
        return f"Error: {str(e)}"
    
def get_version(library_name):
    try:
        module = importlib.import_module(library_name)
        return module.__version__
    except ImportError:
        return "Not installed"
    except AttributeError:
        return "Version not found"

def main():
    libraries = {
        'torch': 'PyTorch',
        'tensorflow': 'TensorFlow',
        'keras': 'Keras',
        'transformers': 'Transformers (Hugging Face)',
        'gym': 'OpenAI Gym',
        'spacy': 'spaCy',
        'nltk': 'NLTK',
        'jax': 'JAX',
        'fastai': 'FastAI'
    }

    print(f"CUDA Version: {get_cuda_version()}")
    print(f"cuDNN Version: {get_cudnn_version()}")
    for lib, name in libraries.items():
        print(f"{name} Version: {get_version(lib)}")

if __name__ == "__main__":
    main()

こちらを適当なファイルで保存して実行すれば結果がprintされます。

gymとかspacyとか不要そうなものも混ざってますが、ここらへんは必要に応じて適宜取り除くなり追加するなりしていただければ大丈夫です。

出力結果

私の環境で動かしてみたところ、こんな感じになりました。

CUDA Version: 12.6, V12.6.77
cuDNN Version: 90100
PyTorch Version: 2.4.0+cu118
TensorFlow Version: Not installed
Keras Version: Not installed
Transformers (Hugging Face) Version: 4.43.0
OpenAI Gym Version: Not installed
spaCy Version: Not installed
NLTK Version: Not installed
JAX Version: 0.4.31
FastAI Version: Not installed

これを見てすぐにCUDAとPyTorchのバージョン思い切り不一致だと気づけた方は流石です（笑）

CUDA関連のバージョン調整をしているところだったので、気にしないでもらえると嬉しいです。