Pytorch导出小样本Prototypical Networks为ONNX时opset9不支持cdist算子报错

当前PyTorch ONNX导出流程默认使用opset 9版本，该版本原生不支持PyTorch提供的cdist距离计算算子，而Prototypical Networks在推理阶段需要调用cdist计算查询样本与各类原型的特征距离，直接触发导出流程的拦截报错。

• 调整导出使用的Opset版本
  ONNX从opset 12开始正式纳入对cdist算子的原生支持，不需要修改模型内部的业务逻辑，仅需在导出代码中指定匹配的opset版本即可完成导出，修改后的示例代码如下：

  torch.onnx.export(
      model,
      (example_query_images, example_query_labels, x_pred),
      "prototypical_network.onnx",
      opset_version=12,  # 从默认的opset9调整为支持cdist的最低兼容版本
      do_constant_folding=True
  )
  

  注意原导出代码中的文件名super_resolution.onnx是超分辨率模型的默认命名，建议替换为对应原型网络的自定义文件名，避免后续使用时出现文件混淆。如果后续要适配的推理引擎对opset版本有上限要求，可以根据引擎的支持范围在12~18的稳定版本区间内选择，不需要盲目使用最高版本。
• 低版本Opset场景下替换等价计算逻辑
  如果必须使用opset9这类老版本导出，无法引入原生cdist支持，可以将模型内部直接调用torch.cdist的代码，替换为数学原理完全等价的展开计算逻辑，从根源上避开不支持的算子。
  替换实现的核心逻辑和原型网络中常用的欧氏距离计算完全对齐，输入输出形状、计算精度和原生cdist(p=2)的表现一致，实现代码如下：

  def euclidean_distance_matrix(query_feat, prototype_feat):
      # query_feat shape: [查询样本数, 特征维度]
      # prototype_feat shape: [类别数, 特征维度]
      query_sq = (query_feat ** 2).sum(dim=1, keepdim=True)
      proto_sq = (prototype_feat ** 2).sum(dim=1, keepdim=True)
      cross_term = query_feat @ prototype_feat.T
      dist_sq = query_sq - 2 * cross_term + proto_sq.T
      # 极小值截断避免开根号出现无意义的数值偏差
      dist_sq = torch.clamp(dist_sq, min=1e-12)
      return torch.sqrt(dist_sq)
  

  直接用这个函数替换模型forward流程里所有的torch.cdist调用，再走原有导出流程即可正常生成ONNX文件，不会触发算子不支持的报错。
• 导出后做精度对齐校验
  不管是选择升级opset版本，还是替换等价计算逻辑，导出完成后都需要拿固定的测试输入，分别运行原生PyTorch模型和导出的ONNX模型，对比两者输出的数值差，误差控制在1e-4以内即说明导出的模型功能正常，没有出现逻辑偏差。

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