怎么自动探查cpu/gpu的 compute_type 并应用呢？

blackantt

compute_type="float16"

ValueError: Requested float16 compute type, but the target device or backend do not support efficient float16 computation.

小甲鱼的二师兄

import torch

def get_compute_type():
    if torch.cuda.is_available():
        # 检查 GPU 是否支持 float16
        if torch.cuda.get_device_capability(0)[0] >= 5.3:
            print("Using float16 precision")
            return "float16"
        else:
            print("float16 not supported on this GPU, using float32")
            return "float32"
    else:
        print("No GPU found, using float32")
        return "float32"

compute_type = get_compute_type()

# 你的模型代码，确保使用合适的精度
if compute_type == "float16":
    model = model.half()
    # 在训练循环中也需要将输入和目标转换为 float16
    # inputs = inputs.half()
    # targets = targets.half()

blackantt

auto

小甲鱼的二师兄

利用深度学习框架（如 TensorFlow、PyTorch）的 API 检查硬件支持情况就可以了。