如何将BERT模型的BaseModelOutputWithPoolingAndCrossAttentions输出转为JSON

要将BaseModelOutputWithPoolingAndCrossAttentions的输出（尤其是last_hidden_state）转换为JSON格式，需先把PyTorch张量转为可JSON序列化的Python数据类型，具体操作如下：

1. 补充必要导入 #

在inference.py中添加json和torch的导入：

import json
import torch
from transformers import BertModel, BertConfig

2. 修改predict_fn处理输出 #

不要直接返回模型原始输出，提取目标张量并做序列化转换：

def predict_fn(input_data, model):
    # 这里是你的输入处理逻辑（比如tokenize）
    inputs = ...
    
    outputs = model(inputs)
    
    # 处理last_hidden_state：移到CPU、移除梯度、转成可序列化的列表
    last_hidden_state = outputs.last_hidden_state.cpu().detach().numpy().tolist()
    
    # 构造结果字典，可按需添加其他字段（如pooler_output）
    result_dict = {
        "last_hidden_state": last_hidden_state,
        "pooler_output": outputs.pooler_output.cpu().detach().numpy().tolist() if outputs.pooler_output is not None else None
    }
    
    # 转成JSON字符串返回
    return json.dumps(result_dict)

关键细节说明 #

• cpu()：如果模型运行在GPU（如示例中的cuda:0），必须把张量移到CPU才能转换为numpy数组。
• detach()：移除张量的梯度信息，带梯度的张量无法直接转numpy，会导致序列化失败。
• .numpy().tolist()：把张量转为numpy数组后再转Python列表，这是JSON支持的标准数据类型。

转换后的JSON示例 #

对应你提供的张量，最终JSON结构大致如下：

{
  "last_hidden_state": [
    [
      [-1.6968, 1.9364, -2.1796, -0.0819, 1.8027, 0.3540, 1.3269, 0.1532],
      [-0.4969, 0.4169, 0.5677, 1.0968, 0.0742, 1.5354, 0.9387, 0.0343]
    ]
  ],
  "pooler_output": null
}

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