PyTorch模型ONNX导出：处理动态控制流与可选输入输出的策略

import torch
import torch.nn as nn

class FormattingLayer(nn.Module):
   def forward(self, input_tensor):
      # 检查输入是否全为零
      is_all_zeros = torch.all(input_tensor == 0)

      # 这里的Python if语句是问题的根源
      if is_all_zeros:
          # 如果输入全为零，期望返回 None
          formatted_input = None
      else:
          # 否则，对输入进行处理
          formatted_input = input_tensor * 2 # 示例处理
      return formatted_input

# 尝试导出模型
model = FormattingLayer()
dummy_input = torch.zeros(2, 3) # 示例输入

# 导出时会收到 Tracer Warning
# torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

import torch
import torch.nn as nn
import torch.jit

class FormattingLayerScripted(nn.Module):
    def forward(self, input_tensor: torch.Tensor):
        # 使用JIT编译，允许内部的Python控制流
        if torch.all(input_tensor == 0):
            # ONNX不支持直接返回 None。
            # 解决方案是返回一个“哨兵值”和/或一个指示是否“忽略”的标志。
            # 例如，返回一个全零张量作为占位符，并返回一个布尔张量表示该输出是否有效。
            placeholder_output = torch.zeros_like(input_tensor)
            is_ignored_flag = torch.tensor(True)
            return placeholder_output, is_ignored_flag
        else:
            # 正常处理逻辑
            formatted_input = input_tensor * 2
            is_ignored_flag = torch.tensor(False)
            return formatted_input, is_ignored_flag

# 实例化并使用 torch.jit.script 编译模型
model_scripted = FormattingLayerScripted()
scripted_model = torch.jit.script(model_scripted)

# 定义示例输入
dummy_input_zeros = torch.zeros(2, 3)
dummy_input_non_zeros = torch.ones(2, 3)

# 导出到ONNX
# 注意：ONNX的输出现在是两个张量：格式化后的张量（或占位符）和表示是否忽略的标志
try:
    torch.onnx.export(scripted_model,
                       dummy_input_zeros, # 使用一个示例输入进行追踪，但JIT已处理了控制流
                       "scripted_model_zeros_input.onnx",
                       input_names=["input_tensor"],
                       output_names=["formatted_output", "is_ignored_flag"],
                       opset_version=11)
    print("Scripted model exported successfully with zeros input.")

    torch.onnx.export(scripted_model,
                       dummy_input_non_zeros, # 另一个示例输入
                       "scripted_model_non_zeros_input.onnx",
                       input_names=["input_tensor"],
                       output_names=["formatted_output", "is_ignored_flag"],
                       opset_version=11)
    print("Scripted model exported successfully with non-zeros input.")

except Exception as e:
    print(f"Error exporting scripted model: {e}")

# 验证输出
# output_zeros, flag_zeros = scripted_model(dummy_input_zeros)
# output_non_zeros, flag_non_zeros = scripted_model(dummy_input_non_zeros)
# print(f"Output for zeros input: {output_zeros}, Ignored: {flag_zeros}")
# print(f"Output for non-zeros input: {output_non_zeros}, Ignored: {flag_non_zeros}")