如何解决Apple M1的MPS设备加载模型时float64转float32的问题

问题原因 #

Apple M1的MPS框架不支持float64（双精度）张量，而你的预训练权重文件中包含这类张量，直接加载到MPS设备就会触发类型转换错误。

解决方案 #

方法一：先加载到CPU，转float32再迁移到MPS

你想到的这个临时方案完全可行，操作步骤如下：

• 先把权重加载到CPU，避免直接触发MPS的类型错误
• 遍历权重字典，把所有张量统一转换为float32
• 最后将整个权重迁移到MPS设备

代码示例：

import torch

# 1. 加载权重到CPU
weights = torch.load('yolov7_training.pt', map_location='cpu')

# 2. 批量转换所有张量为float32
if isinstance(weights, dict):
    for key, tensor in weights.items():
        if isinstance(tensor, torch.Tensor):
            weights[key] = tensor.float()
elif isinstance(weights, torch.Tensor):
    weights = weights.float()

# 3. 迁移到MPS设备
device = torch.device('mps')
weights = weights.to(device)

方法二：用自定义map_location函数，加载时一步转换

通过给torch.load的map_location参数传一个自定义函数，可以在加载每个张量时同时完成设备迁移和类型转换，代码更简洁：

import torch

def convert_to_mps_float32(storage, location):
    # 先加载到CPU转float32，再移到MPS
    return storage.cpu().float().to('mps')

# 加载权重时自动处理类型和设备
weights = torch.load('yolov7_training.pt', map_location=convert_to_mps_float32)

补充说明 #

YOLOv7模型通常是用float32训练的，转换为float32不会影响模型性能，反而能更好适配MPS的硬件加速特性。

内容的提问来源于stack exchange，提问作者Ahasanul Haque