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

import torch
import torch.nn as nn

class FormattingLayer(nn.Module):
    def forward(self, input_tensor):
        # 检查输入是否全为零
        # 原始尝试：torch.gt(torch.nonzero(input_tensor), 0)
        # 更好的检查全零方式：input_tensor.abs().sum() == 0
        is_all_zeros = (input_tensor.abs().sum() == 0)

        if is_all_zeros:
            # 如果全为零，返回 None (原始需求)
            formatted_input = None
        else:
            # 否则，进行格式化处理 (此处简化为原样返回)
            formatted_input = input_tensor # 假设这里有实际的格式化逻辑

        return formatted_input

# 示例模型
model = FormattingLayer()

# 尝试导出为ONNX
dummy_input_zeros = torch.zeros(1, 10)
dummy_input_non_zeros = torch.ones(1, 10)

# 导出全零输入的情况
try:
    torch.onnx.export(model, dummy_input_zeros, "model_zeros.onnx", opset_version=11)
except Exception as e:
    print(f"导出全零输入时出错: {e}")

# 导出非全零输入的情况
try:
    torch.onnx.export(model, dummy_input_non_zeros, "model_non_zeros.onnx", opset_version=11)
except Exception as e:
    print(f"导出非全零输入时出错: {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):
        # 使用张量操作检查是否全为零
        # 注意：TorchScript通常需要将None替换为某种特定值或处理方式
        # ONNX模型输出必须是固定张量，不能是None
        is_all_zeros = (input_tensor.abs().sum() == 0)

        if is_all_zeros:
            # 如果全为零，返回一个全零张量作为“忽略”的信号
            # 原始需求是None，但ONNX不支持None作为输出，需要转换为具体张量
            formatted_input = torch.zeros_like(input_tensor)
        else:
            formatted_input = input_tensor # 实际的格式化逻辑

        return formatted_input

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

# 尝试导出为ONNX
dummy_input_zeros = torch.zeros(1, 10)
dummy_input_non_zeros = torch.ones(1, 10)

# 使用编译后的模型导出
try:
    torch.onnx.export(scripted_model, dummy_input_zeros, "model_scripted_zeros.onnx", opset_version=11)
    print("使用TorchScript成功导出全零输入模型。")
except Exception as e:
    print(f"使用TorchScript导出全零输入模型时出错: {e}")

try:
    torch.onnx.export(scripted_model, dummy_input_non_zeros, "model_scripted_non_zeros.onnx", opset_version=11)
    print("使用TorchScript成功导出非全零输入模型。")
except Exception as e:
    print(f"使用TorchScript导出非全零输入模型时出错: {e}")

import torch
import torch.nn as nn

class FormattingLayerNoControlFlow(nn.Module):
    def forward(self, input_tensor):
        # 1. 检查输入是否全为零
        # input_tensor.abs().sum() > 1e-6 用于判断是否有非零元素
        # 避免使用 == 0，因为浮点数比较可能不精确
        # 结果是一个布尔张量
        has_non_zero_elements = (input_tensor.abs().sum() > 1e-6)

        # 2. 将布尔张量转换为浮点型张量 (0.0 或 1.0)
        # 如果有非零元素，mask为1.0；否则为0.0
        mask = has_non_zero_elements.float()

        # 3. 应用掩码：如果输入被“忽略”，则输出一个全零张量
        # 否则，输出格式化后的输入（此处简化为原样）
        # 这种方式确保输出始终是张量，且形状固定
        formatted_input = input_tensor * mask

        # 或者，如果需要更复杂的条件选择，可以使用torch.where
        # formatted_input = torch.where(has_non_zero_elements, input_tensor, torch.zeros_like(input_tensor))

        return formatted_input

# 实例化模型
model_no_cf = FormattingLayerNoControlFlow()

# 尝试导出为ONNX
dummy_input_zeros = torch.zeros(1, 10)
dummy_input_non_zeros = torch.ones(1, 10)

print("\n--- 尝试导出无Python控制流的模型 ---")
try:
    torch.onnx.export(model_no_cf, dummy_input_zeros, "model_no_cf_zeros.onnx", opset_version=11)
    print("成功导出全零输入模型（无Python控制流）。")
except Exception as e:
    print(f"导出全零输入模型时出错（无Python控制流）: {e}")

try:
    torch.onnx.export(model_no_cf, dummy_input_non_zeros, "model_no_cf_non_zeros.onnx", opset_version=11)
    print("成功导出非全零输入模型（无Python控制流）。")
except Exception as e:
    print(f"导出非全零输入模型时出错（无Python控制流）: {e}")