win、linux等环境下python输出cpu、gpu、avx等硬件信息

测试时输出环境信息，方便判断是否可以开启部分优化

# 环境: 20241030 win10 vs2022 python3.9.13
# vs2022时，在 C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python39_64\Scripts 下运行 pip install 

import os
import sys
import psutil # pip install psutil
import cpuinfo # pip install py-cpuinfo
import subprocess
import platform

# 获取 CPU 插槽数
def get_cpu_socket_count():
    """获取 CPU 插槽数"""
    try:
        if platform.system() == "Windows":
            # 获取 SocketDesignation 字段的输出
            output = subprocess.check_output('wmic cpu get SocketDesignation', shell=True, text=True)
            # 按行分割，排除第一行的标题和空行
            lines = [line.strip() for line in output.splitlines() if line.strip() and "SocketDesignation" not in line]
            return len(lines)  # 返回有效行数作为插槽数
        elif platform.system() == "Linux":
            output = subprocess.check_output('lscpu', shell=True, text=True)
            for line in output.splitlines():
                if "Socket(s):" in line:
                    return int(line.split(":")[1].strip())
    except Exception as e:
        print(f"获取 CPU 插槽数时出错: {e}")
    return "未知"


def get_cpu_model():
    """获取当前 CPU 的型号"""
    try:
        if platform.system() == "Windows":
            output = subprocess.check_output("wmic cpu get Name", shell=True, text=True)
            lines = [line.strip() for line in output.strip().split("\n") if line.strip()]
            if len(lines) > 1:
                return lines[1].strip()  # 返回型号信息
        elif platform.system() == "Linux":
            output = subprocess.check_output("cat /proc/cpuinfo | grep 'model name'", shell=True, text=True)
            for line in output.splitlines():
                if "model name" in line:
                    return line.split(":")[1].strip()
        elif platform.system() == "Darwin":  # macOS
            output = subprocess.check_output("sysctl -n machdep.cpu.brand_string", shell=True, text=True)
            return output.strip()
    except Exception as e:
        print(f"获取 CPU 型号时出错: {e}")
    return "未知型号"

def get_cpu_info():
    """输出运行环境的cpu信息"""
    cpu_model = get_cpu_model()
    socket_count = get_cpu_socket_count()
    physical_cores = psutil.cpu_count(logical=False)  # 物理核心数
    logical_cores = psutil.cpu_count(logical=True)    # 逻辑核心数

    # 获取 CPU 缓存信息
    cpu_info = cpuinfo.get_cpu_info()
    cache_info = {
        'L1': cpu_info.get('l1_cache_size', 'Unknown'),
        'L2': cpu_info.get('l2_cache_size', 'Unknown'),
        'L3': cpu_info.get('l3_cache_size', 'Unknown')
    }

    # 获取所有指令集信息
    cpu_flags = cpu_info.get('flags', [])

    # 检测特定AVX 指令集支持
    avx128_support = 'avx' in cpu_flags
    avx256_support = 'avx2' in cpu_flags
    avx512_support = any(flag.startswith('avx512') for flag in cpu_flags)

    # 输出 CPU 信息
    print("\n--- CPU 信息 ---")
    print(f"CPU 型号: {cpu_model}")
    print(f"CPU 插槽数 (Sockets): {socket_count}")
    print(f"物理核心数 (Physical Cores): {physical_cores}")
    print(f"逻辑核心数 (Logical Cores): {logical_cores}")
    
    # 输出缓存信息
    print("\n--- 缓存信息 ---")
    for level, size in cache_info.items():
        print(f"{level} Cache Size: {size}")
        
    print(f"是否支持 AVX 128 指令集: {'是' if avx128_support else '否'}")
    print(f"是否支持 AVX 256 指令集: {'是' if avx256_support else '否'}")
    print(f"是否支持 AVX 512 指令集: {'是' if avx512_support else '否'}")
    
    # 输出所有支持的指令集
    print("\n支持的指令集 (Flags):")
    for flag in cpu_flags:
        print(f"- {flag}")


# 检查 CPU 信息
get_cpu_info()
# 检查GPU是否可用
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))