M1 Pro运行pc_certification.py遇CUDA及load()错误的解决方法

错误原因及修复方案 #

1. 反序列化错误（CUDA设备不可用） #

M1 Pro芯片使用Apple Silicon的MPS后端而非CUDA，代码中硬编码了CUDA设备，且加载模型时未指定映射到CPU/MPS，导致反序列化失败。

2. TypeError: load()缺失必填位置参数'f' #

通常是模型加载时torch.load调用参数错误，或旧的dump文件与当前设备环境不兼容导致加载失败。

具体代码修改 #

修改主脚本（pc_certification.py） #

（1）动态设置设备

替换硬编码的CUDA设备为自动检测的设备：

# 原代码
device = 'cuda' 
model = model.to(device)
model.eval()
cudnn.benchmark = True

# 修改后
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') 
model = model.to(device)
model.eval()
# 仅在CUDA设备下启用cudnn优化
if device.type == 'cuda':
    cudnn.benchmark = True

（2）替换所有.cuda()调用为.to(device)

找到代码中使用.cuda()的部分，修改为适配当前设备：

# 原代码
masked_output = model(torch.where(torch.logical_and(mask,mask2),data,torch.tensor(0.).cuda()))

# 修改后
masked_output = model(torch.where(torch.logical_and(mask,mask2),data,torch.tensor(0.).to(device)))

修改模型加载工具（utils/setup.py中的get_model函数） #

在加载模型权重时添加map_location参数，确保模型加载到当前设备：

# 原代码（示例）
checkpoint = torch.load(os.path.join(model_dir, f"{model_name}.pth"))

# 修改后
checkpoint = torch.load(os.path.join(model_dir, f"{model_name}.pth"), map_location=device)

如果get_model函数未接收device参数，需要先修改函数定义：

# 修改get_model函数定义
def get_model(model_name, dataset, model_dir, device):
    # ... 原有代码 ...
    checkpoint = torch.load(os.path.join(model_dir, f"{model_name}.pth"), map_location=device)
    # ... 原有代码 ...

并在主脚本中传递device参数：

# 原代码
model = get_model(MODEL_NAME,DATASET,MODEL_DIR)

# 修改后
model = get_model(MODEL_NAME,DATASET,MODEL_DIR, device=device)

重新生成兼容的dump文件 #

执行命令时添加--override参数，强制重新计算并保存适配当前设备的预测结果，避免加载旧的CUDA环境生成的文件：

python pc_certification.py --model resnetv2_50x1_bit_distilled_cutout2_128 --dataset imagenette --num_img -1 --num_mask 6 --patch_size 32 --override

内容的提问来源于stack exchange，提问作者SRAVANTH THATAVARTHI