运行时环境层

运行时环境层

使用 ONNX Runtime 获取可用的执行提供程序，并通过 rocm-smi 工具输出 AMD GPU 具体信息的 Python 代码示例。

• get_amd_gpu_info()：使用 subprocess.run 调用 rocm-smi 工具来获取 AMD GPU 的详细信息。rocm-smi 是 AMD 的工具，用于列出 GPU 的硬件信息。
• check_onnx_runtime_rocm()：调用 ONNX Runtime 的 get_available_providers() 方法，列出 ONNX Runtime 中当前可用的执行提供程序。如果 ROCMExecutionProvider 可用，说明 ROCm 执行提供程序可以在 AMD GPU 上运行。

import subprocess
import onnxruntime as ort

# 获取 AMD GPU 信息
def get_amd_gpu_info():
    try:
        # 使用 rocm-smi 命令获取 GPU 信息
        result = subprocess.run(['rocm-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode == 0:
            print("AMD GPU Info:")
            print(result.stdout)
        else:
            print("Error running rocm-smi:", result.stderr)
    except FileNotFoundError:
        print("rocm-smi not found. Please install ROCm tools.")

# 检查 ONNX Runtime 是否可以使用 ROCm Execution Provider
def check_onnx_runtime_rocm():
    # 获取 ONNX Runtime 中的可用执行提供程序
    providers = ort.get_available_providers()
    print("Available Execution Providers:", providers)
    if 'ROCMExecutionProvider' in providers:
        print("ROCM Execution Provider is available!")
    else:
        print("ROCM Execution Provider is not available.")

if __name__ == "__main__":
    # 输出 AMD GPU 信息
    get_amd_gpu_info()
    
    # 检查 ONNX Runtime 的可用执行提供程序
    check_onnx_runtime_rocm()

结果

AMD GPU Info:


======================================== ROCm System Management Interface ========================================
================================================== Concise Info ==================================================
Device  Node  IDs              Temp    Power  Partitions          SCLK    MCLK   Fan  Perf  PwrCap  VRAM%  GPU%  
              (DID,     GUID)  (Edge)  (Avg)  (Mem, Compute, ID)                                                 
==================================================================================================================
0       1     0x73df,   49956  38.0°C  9.0W   N/A, N/A, 0         500Mhz  96Mhz  0%   auto  203.0W  0%     5%    
==================================================================================================================
============================================== End of ROCm SMI Log ===============================================

Available ONNX Runtime Execution Providers:
TensorrtExecutionProvider
CUDAExecutionProvider
CPUExecutionProvider
ROCm Execution Provider is available.