Pytorch学习笔记

输出关于PyTorch、CUDA设备以及CUDA运行时的相关信息

import torch


def check_torch_and_cuda_details():
    # 检查PyTorch版本
    print("PyTorch version:", torch.__version__)

    # 检查CUDA是否可用
    if torch.cuda.is_available():
        device = torch.device("cuda")

        # 获取CUDA设备属性
        print("\nCUDA Device Properties:")
        print(f"Device Name: {torch.cuda.get_device_name(0)}")
        print(f"Device Capability: {torch.cuda.get_device_capability(0)}")

        # 获取CUDA运行时版本（近似于驱动版本）
        runtime_version = torch.version.cuda
        print(f"CUDA Runtime Version: {runtime_version}")

        # 打印所有设备的信息
        for i in range(torch.cuda.device_count()):
            device_info = torch.cuda.get_device_properties(i)
            print(f"\nDevice {i}:")
            print(f"Name: {device_info.name}")
            print(f"Total Memory (GB): {device_info.total_memory / 1e9:.2f}")
            print(f"Compute Capability: {device_info.major}.{device_info.minor}")

    else:
        print("CUDA is not available")


if __name__ == '__main__':
    check_torch_and_cuda_details()

 运行结果