在ONNX模型中处理PyTorch模块的条件逻辑与可选输入

import torch
import torch.nn as nn

class FormattingLayer(nn.Module):
    def forward(self, input_tensor):
        # 检查输入是否全为零
        # 注意：torch.nonzero(input_tensor).numel() == 0 才是判断全零的可靠方式
        # 原始代码中的 torch.gt(torch.nonzero(input), 0) 可能不准确
        is_all_zeros = (input_tensor == 0).all()

        if is_all_zeros:
            formatted_input = None
        else:
            # 模拟格式化操作
            formatted_input = input_tensor * 2 

        return formatted_input

# 示例模型
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.formatter = FormattingLayer()

    def forward(self, x):
        formatted_x = self.formatter(x)
        # 实际模型中，这里会根据 formatted_x 是否为 None 进行后续处理
        # 例如：if formatted_x is not None: ... else: ...
        return formatted_x

# 尝试转换为ONNX
dummy_input_zeros = torch.zeros(1, 4)
dummy_input_data = torch.randn(1, 4)

model = MyModel()

# 当输入为全零时
try:
    torch.onnx.export(model, dummy_input_zeros, "model_zeros.onnx", 
                      input_names=['input'], output_names=['output'], 
                      opset_version=11, verbose=True)
except Exception as e:
    print(f"\n导出全零输入时遇到错误或警告:\n{e}")

# 当输入包含数据时
try:
    torch.onnx.export(model, dummy_input_data, "model_data.onnx", 
                      input_names=['input'], output_names=['output'], 
                      opset_version=11, verbose=True)
except Exception as e:
    print(f"\n导出包含数据输入时遇到错误或警告:\n{e}")

Tracer Warning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if is_all_zeros:

import torch
import torch.nn as nn

class FormattingLayerScripted(nn.Module):
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        # 使用torch.all()判断是否全为零，这是ONNX兼容的张量操作
        is_all_zeros = (input_tensor == 0).all()

        # 注意：这里不能直接返回 None，因为ONNX模型输出必须是张量
        # 我们返回一个全零张量作为“无输入”的信号
        if is_all_zeros:
            # 返回一个与输入同形的全零张量作为占位符
            # 下游逻辑需要识别这个全零张量为“无有效输入”
            formatted_input = torch.zeros_like(input_tensor)
        else:
            # 模拟格式化操作
            formatted_input = input_tensor * 2 

        return formatted_input

# 使用torch.jit.script编译模块
# 注意：对于复杂的模型，可能需要对整个模型进行脚本化，或者确保子模块是可脚本化的
scripted_formatter = torch.jit.script(FormattingLayerScripted())

# 示例模型
class MyScriptedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.formatter = scripted_formatter # 使用脚本化的模块

    def forward(self, x):
        formatted_x = self.formatter(x)
        # 在ONNX模型中，formatted_x 永远是一个张量
        # 下游逻辑需要根据 formatted_x 的值（例如是否全零）来判断其有效性
        return formatted_x

# 尝试转换为ONNX
dummy_input_zeros = torch.zeros(1, 4)
dummy_input_data = torch.randn(1, 4)

model_scripted = MyScriptedModel()

print("\n--- 导出脚本化模型 (全零输入) ---")
torch.onnx.export(model_scripted, dummy_input_zeros, "model_scripted_zeros.onnx", 
                  input_names=['input'], output_names=['output'], 
                  opset_version=11, verbose=False) # verbose=False 减少输出

print("\n--- 导出脚本化模型 (包含数据输入) ---")
torch.onnx.export(model_scripted, dummy_input_data, "model_scripted_data.onnx", 
                  input_names=['input'], output_names=['output'], 
                  opset_version=11, verbose=False)

print("\n脚本化模型导出成功，且无Tracer Warning。")

import torch
import torch.nn as nn

class FormattingLayerTensorOps(nn.Module):
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        # 判断是否全为零
        is_all_zeros = (input_tensor == 0).all()

        # 将布尔值转换为张量，以便与torch.where 结合
        # 注意：is_all_zeros 是一个0维布尔张量，需要转换为数值类型
        condition_tensor = is_all_zeros.to(torch.bool) # 确保是布尔张量

        # 模拟格式化操作
        formatted_data = input_tensor * 2

        # 使用torch.where根据条件选择输出
        # 如果condition_tensor为True (全零)，则输出全零张量
        # 否则，输出格式化后的数据
        formatted_input = torch.where(condition_tensor, 
                                      torch.zeros_like(input_tensor), 
                                      formatted_data)

        return formatted_input

# 示例模型
class MyTensorOpsModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.formatter = FormattingLayerTensorOps()

    def forward(self, x):
        return self.formatter(x)

# 尝试转换为ONNX
dummy_input_zeros = torch.zeros(1, 4)
dummy_input_data = torch.randn(1, 4)

model_tensor_ops = MyTensorOpsModel()

print("\n--- 导出纯张量操作模型 (全零输入) ---")
torch.onnx.export(model_tensor_ops, dummy_input_zeros, "model_tensor_ops_zeros.onnx", 
                  input_names=['input'], output_names=['output'], 
                  opset_version=11, verbose=False)

print("\n--- 导出纯