Chapter 5: Cross-Framework Conversion

Bridging PyTorch, TensorFlow, and Beyond

弘益人間 · Benefit All Humanity

5.1 The Need for Cross-Framework Conversion

Modern AI development involves multiple frameworks - researchers might use PyTorch for experimentation, while production systems run TensorFlow Serving. Cross-framework conversion enables seamless model transfer, allowing teams to leverage the best tools for each stage of the ML lifecycle.

5.2 Conversion Pathways

ONNX serves as the central hub for most conversions:

Conversion Graph:

PyTorch ←→ ONNX ←→ TensorFlow
   ↓          ↓         ↓
TorchScript  ↓      SavedModel
   ↓          ↓         ↓
Mobile     TFLite   TF.js
   ↓          ↓         ↓
   └──────→ CoreML ←──┘

ONNX → TensorRT (NVIDIA)
ONNX → OpenVINO (Intel)
ONNX → ONNX.js (Web)

5.3 PyTorch to ONNX to TensorFlow

Step 1: PyTorch to ONNX

import torch
import torch.onnx

# PyTorch model
class ResNetBlock(torch.nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(channels, channels, 3, padding=1)
        self.bn1 = torch.nn.BatchNorm2d(channels)
        self.conv2 = torch.nn.Conv2d(channels, channels, 3, padding=1)
        self.bn2 = torch.nn.BatchNorm2d(channels)

    def forward(self, x):
        residual = x
        out = torch.nn.functional.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return torch.nn.functional.relu(out + residual)

model = ResNetBlock(64)
dummy_input = torch.randn(1, 64, 56, 56)

# Export to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet_block.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)

Step 2: ONNX to TensorFlow

import onnx
from onnx_tf.backend import prepare

# Load ONNX model
onnx_model = onnx.load("resnet_block.onnx")

# Convert to TensorFlow
tf_rep = prepare(onnx_model)

# Export as SavedModel
tf_rep.export_graph("resnet_block_tf")

# Or use in Python directly
import numpy as np
output = tf_rep.run(np.random.randn(1, 64, 56, 56).astype(np.float32))
print(output)

5.4 TensorFlow to ONNX to PyTorch

TensorFlow to ONNX

import tensorflow as tf
import tf2onnx

# Create TensorFlow model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64, 3, padding='same', input_shape=(224, 224, 3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.ReLU(),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(1000)
])

# Convert to ONNX
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=17,
    output_path="tf_model.onnx"
)

ONNX to PyTorch

import onnx
import torch
from onnx2pytorch import ConvertModel

# Load ONNX model
onnx_model = onnx.load("tf_model.onnx")

# Convert to PyTorch
pytorch_model = ConvertModel(onnx_model)

# Use the model
test_input = torch.randn(1, 3, 224, 224)
output = pytorch_model(test_input)
print(f"Output shape: {output.shape}")

5.5 Handling Conversion Challenges

Operator Compatibility

Not all operations have direct equivalents across frameworks:

PyTorch	ONNX	TensorFlow	Notes
F.interpolate	Resize	tf.image.resize	Alignment differs
torch.einsum	Einsum	tf.einsum	Opset 12+
torch.nn.GELU	Gelu	tf.nn.gelu	Approximation varies

Addressing Incompatibilities

import torch

# Before conversion: Replace unsupported ops
class ModelWithCompatibleOps(torch.nn.Module):
    def forward(self, x):
        # Instead of: x = torch.special.erf(x)
        # Use supported ops:
        x = 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))  # GELU approx
        return x

# Or use custom ONNX operator
@torch.onnx.symbolic_helper.parse_args('v', 'f')
def custom_op(g, input, scale):
    return g.op("CustomNamespace::CustomOp", input, scale_f=scale)

5.6 Data Layout Differences

PyTorch uses NCHW (batch, channels, height, width) while TensorFlow prefers NHWC:

import torch
import numpy as np

# PyTorch: NCHW
pytorch_tensor = torch.randn(1, 3, 224, 224)

# Convert to TensorFlow: NHWC
tf_tensor = pytorch_tensor.permute(0, 2, 3, 1)  # [1, 224, 224, 3]

# After TensorFlow processing, convert back
pytorch_result = tf_output.permute(0, 3, 1, 2)  # [1, 3, H, W]

# Automatic conversion in ONNX
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output']
)
# ONNX handles layout transformation automatically

5.7 Validation Framework

Always validate converted models to ensure numerical accuracy:

import torch
import numpy as np
import onnxruntime as ort

class ModelValidator:
    def __init__(self, original_model, onnx_path):
        self.pt_model = original_model
        self.ort_session = ort.InferenceSession(onnx_path)

    def validate(self, test_inputs, tolerance=1e-5):
        results = []

        for i, test_input in enumerate(test_inputs):
            # PyTorch inference
            self.pt_model.eval()
            with torch.no_grad():
                pt_output = self.pt_model(test_input).numpy()

            # ONNX inference
            ort_input = {self.ort_session.get_inputs()[0].name: test_input.numpy()}
            ort_output = self.ort_session.run(None, ort_input)[0]

            # Compare
            diff = np.abs(pt_output - ort_output)
            max_diff = diff.max()
            mean_diff = diff.mean()

            passed = max_diff < tolerance

            results.append({
                'test_id': i,
                'passed': passed,
                'max_diff': max_diff,
                'mean_diff': mean_diff
            })

            print(f"Test {i}: {'✓ PASS' if passed else '✗ FAIL'}")
            print(f"  Max diff: {max_diff:.2e}, Mean diff: {mean_diff:.2e}")

        return results

# Usage
validator = ModelValidator(pytorch_model, "model.onnx")
test_data = [torch.randn(1, 3, 224, 224) for _ in range(10)]
results = validator.validate(test_data)

5.8 Converting to Mobile Formats

To TensorFlow Lite

import tensorflow as tf

# From SavedModel
converter = tf.lite.TFLiteConverter.from_saved_model('model_tf')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

# From ONNX (via TensorFlow)
# 1. ONNX → TensorFlow
# 2. TensorFlow → TFLite (as above)

To CoreML

import coremltools as ct

# From ONNX
from onnx_coreml import convert

# Convert
coreml_model = convert(
    model='model.onnx',
    minimum_ios_deployment_target='13',
    preprocessing_args={
        'image_scale': 1.0/255.0,
        'red_bias': -0.485/0.229,
        'green_bias': -0.456/0.224,
        'blue_bias': -0.406/0.225
    }
)

# Save
coreml_model.save('model.mlmodel')

# From PyTorch (via ONNX)
# 1. PyTorch → ONNX
# 2. ONNX → CoreML (as above)

5.9 Automated Conversion Pipelines

from pathlib import Path
import torch
import onnx
import tf2onnx
import tensorflow as tf

class ModelConverter:
    def __init__(self, output_dir="converted_models"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)

    def pytorch_to_all(self, model, dummy_input, name="model"):
        """Convert PyTorch model to multiple formats"""
        results = {}

        # PyTorch → ONNX
        onnx_path = self.output_dir / f"{name}.onnx"
        torch.onnx.export(model, dummy_input, onnx_path, opset_version=17)
        results['onnx'] = onnx_path

        # ONNX → TensorFlow
        try:
            from onnx_tf.backend import prepare
            onnx_model = onnx.load(str(onnx_path))
            tf_path = self.output_dir / f"{name}_tf"
            tf_rep = prepare(onnx_model)
            tf_rep.export_graph(str(tf_path))
            results['tensorflow'] = tf_path
        except Exception as e:
            print(f"TensorFlow conversion failed: {e}")

        # ONNX → TFLite
        try:
            converter = tf.lite.TFLiteConverter.from_saved_model(str(tf_path))
            tflite_model = converter.convert()
            tflite_path = self.output_dir / f"{name}.tflite"
            with open(tflite_path, 'wb') as f:
                f.write(tflite_model)
            results['tflite'] = tflite_path
        except Exception as e:
            print(f"TFLite conversion failed: {e}")

        return results

# Usage
converter = ModelConverter()
results = converter.pytorch_to_all(
    pytorch_model,
    torch.randn(1, 3, 224, 224),
    name="resnet50"
)
print(f"Converted formats: {list(results.keys())}")

5.10 Best Practices