将PyTorch-Lightning模型转为ONNX时遇参数缺失错误求助

TypeError Traceback (most recent call last)
in
22 (input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0)) , export_params=True ,
23 input_names=['images'],
--> 24 output_names=['output'])
25

2 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725 Module: self
726 """
--> 727 return self._apply(lambda t: t.xpu(device))
728
729 def cpu(self: T) -> T:

TypeError: forward() missing 1 required positional argument: 'attention_mask'

模型定义代码如下：
```python
class Tagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.bert = BertModel.from_pretrained(BERT_MODEL_NAME, return_dict=True)
    self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
    self.n_training_steps = n_training_steps
    self.n_warmup_steps = n_warmup_steps
    self.criterion = nn.CrossEntropyLoss(reduction="mean")
    self.n_classes = n_classes
    self.activation = nn.Softmax(dim=1)

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.relu(output)
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True, logger=True)
    return {"loss": loss, "predictions": outputs, "labels": labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True, logger=True)
    return loss

  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True, logger=True)
    return loss

  def training_epoch_end(self, outputs):
    
    labels = []
    predictions = []
    for output in outputs:
      for out_labels in output["labels"].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output["predictions"].detach().cpu():
        predictions.append(out_predictions)

    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)

  def configure_optimizers(self):

    optimizer = AdamW(self.parameters(), lr=2e-5)

    scheduler = get_linear_schedule_with_warmup(
      optimizer,
      num_warmup_steps=self.n_warmup_steps,
      num_training_steps=self.n_training_steps
    )

    return dict(
      optimizer=optimizer,
      lr_scheduler=dict(
        scheduler=scheduler,
        interval='step'
      )
    )

同时希望得到BERT模型生产环境预测的优化思路。

错误原因与修复方案 #

核心问题 #

1. 输入名称不匹配：input_names仅设置了['images']，但模型forward需要两个输入，导致ONNX无法正确映射输入参数
2. 输入维度冗余：tokenizer.encode_plus返回的张量已经是(1, 512)的批量维度，额外的unsqueeze(dim=0)会让维度变成(1,1,512)，不符合BERT的输入要求
3. 输出名称不匹配：模型forward返回(loss, output)两个值，但output_names仅设置了一个，会导致导出异常

修复后的导出代码 #

test_comment = 'I am still waiting on my card?'

encoding = tokenizer.encode_plus(
  test_comment,
  add_special_tokens=True,
  max_length=512,
  return_token_type_ids=False,
  padding="max_length",
  return_attention_mask=True,
  return_tensors='pt',
)

input_ids = encoding["input_ids"]
attention_mask = encoding["attention_mask"]

trained_model.to_onnx(
    "model_lightning_export.onnx",
    (input_ids, attention_mask),  # 移除多余的unsqueeze
    export_params=True,
    input_names=['input_ids', 'attention_mask'],  # 对应两个输入参数
    output_names=['loss', 'output'],  # 对应forward的两个返回值
    dynamic_axes={  # 可选：添加动态维度，支持不同batch size和序列长度的输入
        'input_ids': {0: 'batch_size', 1: 'seq_len'},
        'attention_mask': {0: 'batch_size', 1: 'seq_len'},
        'output': {0: 'batch_size'}
    }
)

推理专用模型导出（可选） #

如果只需要导出用于推理的模型（不需要返回loss），可以封装一个轻量的推理模型：

class InferenceModel(nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model
    
    def forward(self, input_ids, attention_mask):
        _, output = self.model(input_ids, attention_mask)
        return self.model.activation(output)  # 直接返回经过softmax的结果

infer_model = InferenceModel(trained_model).eval()
infer_model.to_onnx(
    "model_lightning_infer.onnx",
    (input_ids, attention_mask),
    export_params=True,
    input_names=['input_ids', 'attention_mask'],
    output_names=['predictions'],
    dynamic_axes={
        'input_ids': {0: 'batch_size', 1: 'seq_len'},
        'attention_mask': {0: 'batch_size', 1: 'seq_len'},
        'predictions': {0: 'batch_size'}
    }
)

BERT模型生产环境预测优化思路 #

• 模型量化：使用PyTorch量化工具或ONNX Runtime将模型从FP32转为INT8，减少内存占用3/4，推理速度提升2-4倍
• 动态批量推理：结合带动态维度的ONNX模型，根据请求量自动调整batch size，避免固定batch size造成的资源浪费
• 模型蒸馏：用大BERT模型蒸馏出DistilBERT或TinyBERT，在精度损失小于5%的前提下，推理速度提升1.5-2倍
• 专用推理引擎：使用ONNX Runtime、TensorRT或VLLM等引擎，相比原生PyTorch推理，速度提升3-10倍
• 输入预处理优化：批量处理输入的tokenization，复用tokenizer实例；缓存高频请求的编码结果
• 部署架构优化：用FastAPI/Flask封装模型为REST API，配合Nginx负载均衡和Kubernetes自动扩缩容，应对高并发
• 模型剪枝：移除冗余的注意力头或神经元，减少模型参数规模，提升推理效率
• 静态图优化：使用TorchScript或ONNX的图优化工具，合并冗余操作，减少推理时的计算量

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