本地与Sagemaker部署sentence-transformers模型推理结果差异求助

本地使用sentence-transformers/distiluse-base-multilingual-cased-v2模型实现文本转向量，代码如下：

class SequenceEncoder(object):
    def __init__(self, device=None):
        self.device = device
        self.multi_model = SentenceTransformer('distiluse-base-multilingual-cased-v2',device=device)   # 512 dimensional dense vector

    @torch.no_grad()
    def __call__(self, col, dialect_list=None):
        if isinstance(col, str):
            vals = [col]
        else:
            vals = col.replace([np.nan, 0], '').values.tolist()
        x = self.multi_model.encode(vals, show_progress_bar=True, convert_to_tensor=True, device=self.device)
        return x.cpu()

本地输出示例：[-0.00285156 -0.04651115 -0.00723144 -0.04229123 -0.02418377,0.00646215, ...]

将同款模型部署到Amazon Sagemaker，配置参数：

• HF_MODEL_ID: sentence-transformers/distiluse-base-multilingual-cased-v2
• HF_TASK: feature-extraction

部署成功后用相同文本测试，发现输出与本地存在明显差异：
Sagemaker输出示例：[-0.035367466509342194, 0.011641714721918106, -0.04396483674645424, 0.03655952587723732, ...]

核心差异点：

1. 数值差异：Sagemaker返回的向量数值与本地模型输出完全不同
2. 浮点精度差异：Sagemaker返回的浮点数值保留位数更长

已尝试移除本地模型调用时的convert_to_tensor=True参数，问题未解决。

需求：

• 明确差异产生的原因
• 实现让Sagemaker返回张量格式的输出

差异产生的原因 #

1. 推理逻辑不匹配 #

本地使用SentenceTransformer的encode方法时，内部会自动完成文本预处理、模型前向传播、句子级池化整套流程，最终输出的是句子级的语义向量。而Sagemaker默认的feature-extraction任务，只会返回模型最后一层的token级隐藏状态，不会执行SentenceTransformer特有的池化逻辑，这是数值差异的核心原因。

2. 数据类型转换导致精度显示差异 #

本地返回的是PyTorch张量默认的float32精度，打印时会自动截断小数位；而Sagemaker推理容器会默认将张量转换为Python原生浮点数返回，保留了更多小数位，属于显示层面的精度差异，而非实际计算精度问题。

3. 环境与配置偏差 #

• 本地与Sagemaker可能使用不同版本的PyTorch/Transformers库，底层实现细节的差异会影响输出数值
• Sagemaker容器可能默认启用了混合精度推理，进一步导致数值偏差

解决方法 #

1. 统一推理逻辑，消除数值差异 #

不要使用默认的feature-extraction任务，自定义推理脚本复现SentenceTransformer.encode的完整逻辑：

创建inference.py脚本：

from sentence_transformers import SentenceTransformer
import torch

def model_fn(model_dir):
    model = SentenceTransformer(model_dir)
    return model

def predict_fn(input_data, model):
    with torch.no_grad():
        outputs = model.encode(input_data, convert_to_tensor=True)
    return outputs.cpu().numpy().tolist()

部署时指定该脚本，确保Sagemaker执行和本地完全一致的池化流程，数值差异即可消除。

2. 让Sagemaker返回张量格式 #

Sagemaker默认返回Python列表，若要直接返回张量，可在自定义脚本中将张量序列化为字节形式返回，本地再反序列化：

修改后的predict_fn：

import pickle

def predict_fn(input_data, model):
    with torch.no_grad():
        outputs = model.encode(input_data, convert_to_tensor=True)
    return pickle.dumps(outputs.cpu())

本地调用Sagemaker端点后，使用pickle.loads()即可反序列化得到PyTorch张量。

3. 统一浮点精度显示 #

若需要统一本地与Sagemaker的浮点显示精度，可在自定义脚本中对输出做精度截断：

def predict_fn(input_data, model):
    with torch.no_grad():
        outputs = model.encode(input_data, convert_to_tensor=True)
    return outputs.cpu().float().numpy().tolist()

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