Azure ML部署MLflow模型时，如何设置map_location=torch.device('cpu')？

针对你遇到的RuntimeError（在无CUDA环境下尝试反序列化CUDA设备上的模型），以下是几种无需修改Docker环境文件的解决方案：

方法一：改用MLflow PyTorch专用加载方法 #

mlflow.pytorch.load_model 原生支持map_location参数，可直接指定加载到CPU，若需兼容PyFunc接口，可转换为PyFunc模型：

import os
import mlflow.pytorch
import torch

model_path = os.path.join(os.getenv("AZUREML_MODEL_DIR"), "use-case1-model")
# 直接指定map_location到CPU
pytorch_model = mlflow.pytorch.load_model(model_path, map_location=torch.device('cpu'))
# 如需PyFunc接口，转换模型
pyfunc_model = mlflow.pyfunc.PyFuncModel(model_impl=pytorch_model)

方法二：通过环境变量强制加载到CPU #

在Azure ML在线端点的部署配置中添加以下环境变量之一，无需修改代码即可让模型自动加载到CPU：

• 设置 MLFLOW_PYTORCH_LOAD_MODEL_MAP_LOCATION=cpu：MLflow专属环境变量，直接控制PyTorch模型的加载位置
• 设置 CUDA_VISIBLE_DEVICES=""：让PyTorch检测不到GPU，自动 fallback 到CPU
• 设置 TORCH_DEVICE=cpu：指定PyTorch默认设备为CPU

方法三：Monkey Patch torch.load 强制添加map_location #

如果必须使用mlflow.pyfunc.load_model，可在score.py开头修改torch.load的默认行为，自动添加map_location参数：

import torch
# 保存原始torch.load方法
original_torch_load = torch.load

def patched_torch_load(*args, **kwargs):
    # 自动添加map_location=cpu，不覆盖已有设置
    kwargs.setdefault('map_location', torch.device('cpu'))
    return original_torch_load(*args, **kwargs)

# 替换torch.load
torch.load = patched_torch_load

# 正常加载PyFunc模型
import os
import mlflow.pyfunc

model_path = os.path.join(os.getenv("AZUREML_MODEL_DIR"), "use-case1-model")
model = mlflow.pyfunc.load_model(model_path)

内容的提问来源于stack exchange，提问作者Rishav Raj