在GPU实例部署端点时如何将模型加载至CUDA设备

数据老陈

2026-8-13

解决PyTorch模型在SageMaker GPU实例加载到CPU的问题

针对你在ml.g4dn.xlarge实例上部署PyTorch模型却加载到CPU的问题，可按以下步骤排查解决：

1. 确保推理代码将模型迁移到CUDA设备

你当前的代码仅定义了DEVICE变量，但未将加载后的模型移动到目标设备上。这是最常见的原因，修改推理代码如下：

import torch
import logging

logger = logging.getLogger(__name__)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
logger.info('Model will be loaded into:{}'.format(DEVICE))

# 加载模型（替换为你的实际模型加载逻辑）
model = torch.load('model.pth')
# 关键步骤：将模型迁移到指定设备
model = model.to(DEVICE)
# 设置模型为推理模式
model.eval()

def predict(input_data):
    # 将输入数据同步迁移到CUDA设备
    inputs = torch.tensor(input_data).to(DEVICE)
    with torch.no_grad():
        outputs = model(inputs)
    # 若需返回CPU格式结果，再迁移回CPU
    return outputs.cpu().numpy()

2. 显式配置SageMaker模型的GPU环境

在创建PyTorchModel时，添加环境变量明确启用GPU支持：

model = PyTorchModel(
    entry_point='inference.py',
    source_dir='code',
    role=role,
    model_data=model_data,
    framework_version="1.12.0",
    py_version='py38',
    code_location='s3://staging',
    name='Staging-Model',
    # 添加GPU启用环境变量
    environment={"SAGEMAKER_USE_GPU": "true"}
)

3. 添加CUDA检测日志排查问题

在推理代码中增加更详细的CUDA状态日志，确认Torch是否真的检测到GPU：

logger.info(f"CUDA available: {torch.cuda.is_available()}")
logger.info(f"GPU数量: {torch.cuda.device_count()}")
if torch.cuda.is_available():
    logger.info(f"当前GPU索引: {torch.cuda.current_device()}")
    logger.info(f"GPU名称: {torch.cuda.get_device_name(0)}")

通过这些日志可以区分是Torch未检测到GPU，还是代码未正确迁移模型。