2.3.5

2.3.5

Important

Starting with onnx2tf v2.4.0, tf_converter will be deprecated and the default backend will be switched to flatbuffer_direct. With the v2.3.3 update, all backward compatible conversion options have been migrated to flatbuffer_direct, so I will only be doing minor bug fixes until April. If you provide us with ONNX sample models, I will consider incorporating them into flatbuffer_direct.

Summary

This PR significantly expands the feat-torch work around the flatbuffer_direct pipeline, with a primary focus on native PyTorch package generation, direct evaluation, and end-to-end validation.

The branch moves PyTorch export from an experimental wrapper-heavy path toward a practical direct-codegen flow built on ModelIR, with better output readability, stronger validation, and broader compatibility for both ONNX-input and TFLite-input workflows.

What Changed

1. Added native PyTorch package generation from ModelIR

Added a path to generate PyTorch model source code from a LiteRT model without going through ONNX. This is a very experimental implementation and has only been tested on a few models.

• Introduced direct PyTorch package export from flatbuffer_direct ModelIR.
• Added a dedicated PyTorch exporter and runtime:
  • onnx2tf/tflite_builder/pytorch_exporter.py
  • onnx2tf/tflite_builder/pytorch_package_runtime.py
• Native packages now generate a real torch.nn.Module-based model.py instead of always falling back to a wrapper package.
• Generated native packages support:
  • Model(...)
  • load_model(...)
  • state_dict.pth loading through standard load_state_dict(...)
• The generated state_dict.pth format was updated to be compatible with standard PyTorch state_dict() keys for native packages.

A sample PyTorch model model.py generated from LiteRT

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable, Dict, Optional

import torch
import torch.nn.functional as F

from .runtime import (
    _align_tensor_to_target_shape,
    _apply_module_transpose_conv2d,
    _normalize_axes,
    _reduce_mean,
    _resolve_named_input_value,
    _shape_list,
    load_generated_weights,
)

PACKAGE_DIR = Path(__file__).resolve().parent
INPUT_NAMES = ['onnx____Reshape_0']
OUTPUT_NAMES = ['63']
class _Conv2dBlock(torch.nn.Module):
    def __init__(
        self,
        *,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple[int, int],
        stride: tuple[int, int],
        padding: tuple[int, int],
        dilation: tuple[int, int],
        groups: int,
        bias: bool,
        pad: Optional[list[int]] = None,
        activation: str = 'none',
        negative_slope: float = 0.2,
        pad_mode: str = 'constant',
        pad_value: float = 0.0,
    ) -> None:
        super().__init__()
        self.conv = torch.nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )
        self.pad = pad
        self.activation = str(activation)
        self.negative_slope = float(negative_slope)
        self.pad_mode = str(pad_mode)
        self.pad_value = float(pad_value)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.pad is not None:
            x = F.pad(x, self.pad, mode=self.pad_mode, value=self.pad_value)
        x = self.conv(x)
        if self.activation == 'leaky_relu':
            return F.leaky_relu(x, negative_slope=self.negative_slope)
        if self.activation == 'relu':
            return torch.relu(x)
        if self.activation == 'relu6':
            return torch.clamp(x, min=0.0, max=6.0)
        if self.activation == 'relu_n1_to_1':
            return torch.clamp(x, min=-1.0, max=1.0)
        if self.activation == 'relu_0_to_1':
            return torch.clamp(x, min=0.0, max=1.0)
        if self.activation == 'tanh':
            return torch.tanh(x)
        if self.activation == 'sigmoid':
            return torch.sigmoid(x)
        return x

class Model(torch.nn.Module):
    const_ConvTranspose_12_transpose_conv_bias: torch.Tensor
    const_ConvTranspose_15_transpose_conv_bias: torch.Tensor
    const_ConvTranspose_18_transpose_conv_bias: torch.Tensor
    const_decoder_decoder_lin_0_bias: torch.Tensor
    const_onnx__MatMul_73: torch.Tensor
    const_onnx__MatMul_74: torch.Tensor
    def __init__(self, *, device: str | None = None, eval_mode: bool = True, load_weights: bool = True):
        super().__init__()
        self.input_names = list(INPUT_NAMES)
        self.output_names = list(OUTPUT_NAMES)
        self.conv_block_0 = _Conv2dBlock(
            in_channels=1,
            out_channels=64,
            kernel_size=(1, 9),
            stride=(1, 1),
            padding=(0, 4),
            dilation=(1, 1),
            groups=1,
            bias=True,
            pad=None,
            activation='relu',
            negative_slope=0.2,
            pad_mode='constant',
            pad_value=0.0,
        )
        self.conv_block_1 = _Conv2dBlock(
            in_channels=64,
            out_channels=128,
            kernel_size=(1, 5),
            stride=(1, 1),
            padding=(0, 2),
            dilation=(1, 1),
            groups=1,
            bias=True,
            pad=None,
            activation='relu',
            negative_slope=0.2,
            pad_mode='constant',
            pad_value=0.0,
        )
        self.conv_block_2 = _Conv2dBlock(
            in_channels=128,
            out_channels=256,
            kernel_size=(1, 3),
            stride=(1, 1),
            padding=(0, 0),
            dilation=(1, 1),
            groups=1,
            bias=True,
            pad=[2, 2, 0, 0, 0, 0, 0, 0],
            activation='none',
            negative_slope=0.2,
            pad_mode='constant',
            pad_value=0.0,
        )
        self.conv_transpose2d_0 = torch.nn.ConvTranspose2d(
            in_channels=256,
            out_channels=128,
            kernel_size=(1, 3),
            stride=(1, 1),
            padding=(0, 0),
            bias=False,
        )
        self.conv_transpose2d_1 = torch.nn.ConvTranspose2d(
            in_channels=128,
            out_channels=64,
            kernel_size=(1, 5),
            stride=(1, 1),
            padding=(0, 0),
            bias=False,
        )
        self.conv_transpose2d_2 = torch.nn.ConvTranspose2d(
            in_channels=64,
            out_channels=64,
            kernel_size=(1, 9),
            stride=(1, 1),
            padding=(0, 0),
            bias=False,
        )
        self._init_constants()
        if load_weights:
            load_generated_weights(
                model=self,
                package_dir=PACKAGE_DIR,
                device=device,
            )
        elif device is not None:
            self.to(device)
        if eval_mode:
            self.eval()

    def _init_constants(self) -> None:
        self.register_buffer('const_ConvTranspose_12_transpose_conv_bias', torch.zeros([1, 1, 1, 128], dtype=torch.float32), persistent=True)
        self.register_buffer('const_ConvTranspose_15_transpose_conv_bias', torch.zeros([1, 1, 1, 64], dtype=torch.float32), persistent=True)
        self.register_buffer('const_ConvTranspose_18_transpose_conv_bias', torch.zeros([1, 1, 1, 64], dtype=torch.float32), persistent=True)
        self.register_buffer('const_decoder_decoder_lin_0_bias', torch.zeros([1, 66, 1], dtype=torch.float32), persistent=True)
        self.register_buffer('const_onnx__MatMul_73', torch.zeros([66, 2], dtype=torch.float32), persistent=True)
        self.register_buffer('const_onnx__MatMul_74', torch.zeros([2, 66], dtype=torch.float32), persistent=True)

    def _device(self) -> torch.device:
        for parameter in self.parameters():
            return parameter.device
        for buffer in self.buffers():
            return buffer.device
        return torch.device('cpu')

    def _max_pool2d_same(self, x: torch.Tensor, *, kernel_size: tuple[int, int], stride: tuple[int, int]) -> torch.Tensor:
        pad_h_total = max(int(kernel_size[0]) - int(stride[0]), 0)
        pad_w_total = max(int(kernel_size[1]) - int(stride[1]), 0)
        pad_top = pad_h_total // 2
        pad_bottom = pad_h_total - pad_top
        pad_left = pad_w_total // 2
        pad_right = pad_w_total - pad_left
        x = F.pad(x, [pad_left, pad_right, pad_top, pad_bottom], mode='constant', value=float('-inf'))
        return F.max_pool2d(x, kernel_size=kernel_size, stride=stride)

    def forward(self, *args: torch.Tensor, **kwargs: torch.Tensor) -> Any:
        if len(args) > 0 and len(kwargs) > 0:
            raise RuntimeError('Use either positional inputs or keyword inputs, not both.')
        if len(kwargs) > 0:
            onnx____Reshape_0 = _resolve_named_input_value(kwargs, 'onnx____Reshape_0')
        else:
            if len(args) != 1:
                raise RuntimeError(f'Input arity mismatch. expected={len(self.input_names)} actual={len(args)}')
            onnx____Reshape_0 = args[0]
        input = torch.reshape(onnx____Reshape_0, [int(v) for v in _shape_list(torch.as_tensor([-1, 1, 64], dtype=torch.int32, device=self._device()))])
        Conv_2_conv1d_input_nchw2d = torch.reshape(input, [int(v) for v in _shape_list(torch.as_tensor([1, 1, 1, 64], dtype=torch.int32, device=self._device()))])
        Conv_2_input = torch.reshape(Conv_2_conv1d_input_nchw2d, [int(v) for v in _shape_list(torch.as_tensor([1, 1, 64, 1], dtype=torch.int32, device=self._device()))])
        Conv_4_input = self.conv_block_0(Conv_2_input.permute(0, 1, 3, 2).contiguous())
        Conv_6_input = self.conv_block_1(Conv_4_input)
        Conv_6_output = self.conv_block_2(Conv_6_input)
        input_31 = torch.reshape(Conv_6_output.permute(0, 3, 1, 2).contiguous(), [int(v) for v in [1, 66, 256]])
        input_35 = torch.relu(input_31)
        _tmp_x_11 = input_35
        _tmp_y_11 = self.const_onnx__MatMul_73
        _tmp_x_11 = _tmp_x_11.transpose(-1, -2)
        onnx____Add_49 = _align_tensor_to_target_shape(torch.matmul(_tmp_x_11, _tmp_y_11), [1, 256, 2])
        onnx____MatMul_50 = torch.add(torch.as_tensor([0.05643776059150696, -0.00908423587679863], dtype=torch.float32, device=self._device()).reshape([1, 1, 2]), onnx____Add_49)
        _tmp_x_13 = self.const_onnx__MatMul_74
        _tmp_y_13 = onnx____MatMul_50
        _tmp_x_13 = _tmp_x_13.transpose(-1, -2)
        _tmp_y_13 = _tmp_y_13.transpose(-1, -2)
        onnx____Add_52 = torch.matmul(_tmp_x_13, _tmp_y_13)
        onnx____ConvTranspose_53 = torch.add(self.const_decoder_decoder_lin_0_bias, onnx____Add_52)
        ConvTranspose_12_convtranspose1d_input_nchw2d = torch.reshape(onnx____ConvTranspose_53, [int(v) for v in _shape_list(torch.as_tensor([1, 1, 66, 256], dtype=torch.int32, device=self._device()))])
        ConvTranspose_12_output = _apply_module_transpose_conv2d(self.conv_transpose2d_0, ConvTranspose_12_convtranspose1d_input_nchw2d, target_shape=[1, 1, 68, 128], fallback_shape=[1, 1, 68, 128], fused='NONE')
        ConvTranspose_12_output_nhwc_cropped = ConvTranspose_12_output[0:1, 0:1, 2:66, 0:128]
        ConvTranspose_12_output_nhwc_bias = torch.add(ConvTranspose_12_output_nhwc_cropped, self.const_ConvTranspose_12_transpose_conv_bias)
        ConvTranspose_15_input = torch.relu(ConvTranspose_12_output_nhwc_bias)
        ConvTranspose_15_output = _apply_module_transpose_conv2d(self.conv_transpose2d_1, ConvTranspose_15_input, target_shape=[1, 1, 68, 64], fallback_shape=[1, 1, 68, 64], fused='NONE')
        ConvTranspose_15_output_nhwc_cropped = ConvTranspose_15_output[0:1, 0:1, 2:66, 0:64]
        ConvTranspose_15_output_nhwc_bias = torch.add(ConvTranspose_15_output_nhwc_cropped, self.const_ConvTranspose_15_transpose_conv_bias)
        ConvTranspose_18_input = torch.relu(ConvTranspose_15_output_nhwc_bias)
        ConvTranspose_18_output = _apply_module_transpose_conv2d(self.conv_transpose2d_2, ConvTranspose_18_input, target_shape=[1, 1, 72, 64], fallback_shape=[1, 1, 72, 64], fused='NONE')
        ConvTranspose_18_output_nhwc_cropped = ConvTranspose_18_output[0:1, 0:1, 4:68, 0:64]
        ConvTranspose_18_output_nhwc_bias = torch.add(ConvTranspose_18_output_nhwc_cropped, self.const_ConvTranspose_18_transpose_conv_bias)
        onnx____ReduceMean_60 = torch.reshape(ConvTranspose_18_output_nhwc_bias, [int(v) for v in _shape_list(torch.as_tensor([1, 64, 64], dtype=torch.int32, device=self._device()))])
        onnx____Squeeze_61 = _reduce_mean(onnx____ReduceMean_60, _normalize_axes(torch.as_tensor([2], dtype=torch.int32, device=self._device()), onnx____ReduceMean_60.ndim), True)
        t_63 = torch.reshape(onnx____Squeeze_61, [int(v) for v in _shape_list(torch.as_tensor([1, 64], dtype=torch.int32, device=self._device()))])
        return t_63

    def forward_named(self, *args: torch.Tensor, **kwargs: torch.Tensor) -> Dict[str, torch.Tensor]:
        return {'63': self.forward(*args, **kwargs)}

def load_model(device: str | None = None, eval_mode: bool = True) -> Model:
    return Model(device=device, eval_mode=eval_mode)

2. Made generated model.py more PyTorch-like

• Refactored code generation to favor a PyTorch-first style instead of exposing low-level export/runtime details.
• Simplified generated source for native packages by:
  • reducing wrapper/template output
  • moving low-level loading/runtime helpers out of model.py
  • preferring direct torch / torch.nn.functional expressions
  • introducing block-style generation for repeated patterns
  • splitting complex forward() methods into stage/helper structure where appropriate
• Improved readability for simple models such as fused_conv_relu.
• Extended the same philosophy to more complex models such as:
  • yolov7_tiny_head_0.768_post_480x640
  • yolox_s
  • lite_model_mobilebert_1_metadata_1

3. Expanded native operator/codegen coverage

• Added or improved native export/codegen/runtime support for a broad set of ops and graph patterns, including coverage needed by real models in target_models.txt.
• This includes work across layout propagation, recurrent handling, detection postprocessing, reshape/layout synchronization, and direct lowering for operators such as:
  • GATHER, GATHER_ND
  • SLICE, STRIDED_SLICE
  • RESIZE_*
  • LOGISTIC, SOFTMAX
  • DEPTH_TO_SPACE, SPACE_TO_DEPTH
  • ARG_MAX, ARG_MIN
  • AVERAGE_POOL_2D, MAX_POOL_2D
  • NON_MAX_SUPPRESSION_V4/V5
  • recurrent sequence ops and related state-shape handling
• Strengthened channel-first normalization and layout propagation so native export can stay closer to optimized PyTorch execution instead of bouncing through stale NHWC-style intermediates.

4. Added direct PyTorch accuracy evaluation and comparison reports

• Added PyTorch package evaluation support:
  • onnx2tf/tflite_builder/pytorch_accuracy_evaluator.py
• -tb flatbuffer_direct -fdopt -cotof now supports same-input comparison reporting for:
  • ONNX ↔ TFLite
  • ONNX ↔ PyTorch
• Added combined comparison reports that summarize both results side-by-side.
• Introduced task-aware pass handling for cases such as probability-map outputs where exact tensor equality is not the right semantic criterion.
• Added handling for ONNX Runtime inference failures so that PyTorch validation can continue when the user explicitly wants those exceptions ignored.

5. Added official -it + -fdopt support

• Formalized support for using --input_tflite_file_path together with:
  • -tb flatbuffer_direct
  • -fdopt
  • -cotof
• For -it + -fdopt + -cotof, the reference comparison is now TFLite ↔ PyTorch instead of ONNX ↔ PyTorch.
• Added report generation for direct TFLite-input workflows.
• Updated CLI behavior and README examples accordingly.

6. Added PyTorch bulk verification tooling

• Added a reusable bulk verification CLI:
  • onnx2tf-pytorch-bulk
  • implemented in onnx2tf/utils/pytorch_bulk_runner.py
• The bulk runner supports:
  • ordered processing of target_models.txt
  • duplicate preservation
  • skipping missing models
  • resume support
  • per-model logs and reports
  • PyTorch package smoke/accuracy evaluation
• This was used to drive real-model closure work across many ONNX models.

7. Improved flatbuffer_direct SavedModel/TFLite bridge behavior

• Added/expanded direct-export and bridge behavior for SavedModel and PyTorch outputs.
• Avoids materializing unnecessary hidden SavedModel bridge directories when native PyTorch export succeeds.
• Added missing SavedModel exporter kernels such as:
  • RANDOM_UNIFORM
  • REVERSE_SEQUENCE
• Kept flatbuffer_direct validation and report generation aligned across ONNX-input and TFLite-input paths.

8. Documentation and test coverage

• Updated README.md extensively for the new PyTorch export flow.
• Added examples for:
  • flatbuffer_direct PyTorch export
  • TFLite-input direct PyTorch export (-it + -fdopt)
  • comparison reporting with -cotof
• Added explicit documentation that generated native PyTorch packages are inference-oriented and should not be treated as training-ready models without further redesign.
• Added or expanded large test coverage across:
  • native codegen
  • accuracy evaluators
  • bulk runner
  • split/export flows
  • direct TFLite-input workflows

Why This Matters

This branch makes flatbuffer_direct much more practical as a direct export pipeline rather than only a TFLite writer. In particular, it enables:

• native PyTorch package output from direct-exported graphs
• direct validation against ONNX or TFLite references
• better generated PyTorch source quality
• larger real-model coverage
• direct TFLite-input to PyTorch workflows without falling back to the legacy converter path

Testing

Verified with focused tests during implementation and with the flatbuffer_direct suite:

pytest -q tests/test_tflite2sm_phase1.py tests/test_pytorch_exporter.py tests/test_model_convert.py tests/test_flatbuffer_direct_op_error_report.py tests/test_tflite_builder_direct.py

Result:

• 793 passed
• 4 warnings

Notes

• The generated native PyTorch packages are intended primarily for inference/package validation, not as drop-in training architectures.
• Existing public dependency constraints were not changed in pyproject.toml; this branch only bumps the package version.

What's Changed

• Improve flatbuffer_direct PyTorch export and validation workflow by @PINTO0309 in #905

Full Changelog: 2.3.4...2.3.5