Chapter 4: Cross-Framework Conversion

WIA-AI-008 Standard • Estimated reading time: 60 minutes

4.1 The Cross-Framework Challenge

Modern AI development often requires converting models between frameworks. A research team might train in PyTorch but need to deploy with TensorFlow Serving. A mobile app might require converting a TensorFlow model to CoreML for iOS or TFLite for Android. Understanding conversion paths and potential pitfalls is essential for successful model exchange.

Common Conversion Paths

Source	Target	Method	Difficulty
PyTorch	ONNX	torch.onnx.export()	Easy
TensorFlow	ONNX	tf2onnx	Easy
ONNX	TensorFlow	onnx-tf	Medium
PyTorch	TFLite	PyTorch→ONNX→TF→TFLite	Hard
TensorFlow	CoreML	coremltools	Medium

4.2 PyTorch to TensorFlow

Converting from PyTorch to TensorFlow typically uses ONNX as an intermediate format:

Step-by-Step Conversion

# Step 1: PyTorch to ONNX
import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    opset_version=13,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)

# Step 2: ONNX to TensorFlow
from onnx_tf.backend import prepare
import onnx

onnx_model = onnx.load("resnet18.onnx")
tf_model = prepare(onnx_model)

# Export as TensorFlow SavedModel
tf_model.export_graph("resnet18_tf")

# Step 3: Load and use in TensorFlow
import tensorflow as tf

loaded_model = tf.saved_model.load("resnet18_tf")
infer = loaded_model.signatures["serving_default"]

# Test inference
import numpy as np
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = infer(input=tf.constant(test_input))

print(f"Output shape: {output['output'].shape}")

Validation After Conversion

import torch
import tensorflow as tf
import numpy as np

# Prepare test input
test_input_np = np.random.randn(5, 3, 224, 224).astype(np.float32)

# PyTorch inference
model_pt = models.resnet18(pretrained=True)
model_pt.eval()
with torch.no_grad():
    output_pt = model_pt(torch.from_numpy(test_input_np)).numpy()

# TensorFlow inference
model_tf = tf.saved_model.load("resnet18_tf")
infer_tf = model_tf.signatures["serving_default"]
output_tf = infer_tf(input=tf.constant(test_input_np))['output'].numpy()

# Compare outputs
max_diff = np.max(np.abs(output_pt - output_tf))
mean_diff = np.mean(np.abs(output_pt - output_tf))
print(f"Max difference: {max_diff:.6f}")
print(f"Mean difference: {mean_diff:.6f}")

assert max_diff < 1e-4, "Conversion validation failed!"
print("✓ Conversion validated successfully")

4.3 TensorFlow to PyTorch

The reverse conversion also uses ONNX:

# Step 1: TensorFlow to ONNX
import tensorflow as tf
import tf2onnx

# Load TensorFlow model
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 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=13)

with open("mobilenet_v2.onnx", "wb") as f:
    f.write(model_proto.SerializeToString())

# Step 2: ONNX to PyTorch
import onnx
from onnx2pytorch import ConvertModel

onnx_model = onnx.load("mobilenet_v2.onnx")
pytorch_model = ConvertModel(onnx_model)

# Step 3: Use in PyTorch
import torch

pytorch_model.eval()
test_input = torch.randn(1, 3, 224, 224)

with torch.no_grad():
    output = pytorch_model(test_input)

print(f"Output shape: {output.shape}")

4.4 Converting to Mobile Formats

PyTorch to TensorFlow Lite

For Android and embedded systems:

# Complete pipeline: PyTorch → ONNX → TF → TFLite
import torch
import onnx
from onnx_tf.backend import prepare
import tensorflow as tf

# 1. PyTorch to ONNX
model_pt = torch.nn.Sequential(
    torch.nn.Conv2d(3, 32, 3, padding=1),
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(2),
    torch.nn.Flatten(),
    torch.nn.Linear(32 * 112 * 112, 10)
)

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model_pt, dummy_input, "model.onnx", opset_version=13)

# 2. ONNX to TensorFlow
onnx_model = onnx.load("model.onnx")
tf_model = prepare(onnx_model)
tf_model.export_graph("model_tf")

# 3. TensorFlow to TFLite
converter = tf.lite.TFLiteConverter.from_saved_model("model_tf")

# Apply optimizations
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()

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

print(f"TFLite model size: {len(tflite_model) / 1024:.2f} KB")

TensorFlow to CoreML

For iOS deployment:

import coremltools as ct
import tensorflow as tf

# Load TensorFlow model
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# Convert to CoreML
coreml_model = ct.convert(
    model,
    inputs=[ct.ImageType(name="input", shape=(1, 224, 224, 3))],
    classifier_config=ct.ClassifierConfig(
        class_labels=['class1', 'class2', ...]  # Your class labels
    )
)

# Add metadata
coreml_model.author = 'WIA Team'
coreml_model.license = 'Apache 2.0'
coreml_model.short_description = 'MobileNetV2 image classifier'
coreml_model.version = '1.0.0'

# Save
coreml_model.save("MobileNetV2.mlmodel")

print("✓ CoreML model created successfully")

4.5 Handling Architecture Differences

Batch Normalization

PyTorch and TensorFlow implement BatchNorm differently:

# PyTorch: momentum = 0.1 (default)
# TensorFlow: momentum = 0.99 (default)

# PyTorch BatchNorm
bn_pt = torch.nn.BatchNorm2d(64, momentum=0.1)

# Equivalent TensorFlow BatchNorm
bn_tf = tf.keras.layers.BatchNormalization(momentum=0.9)  # Note: 1 - 0.1 = 0.9

# When converting, ensure momentum values are adjusted

Padding Differences

# PyTorch: 'same' padding not available in Conv2d (use padding parameter)
conv_pt = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)  # Manual padding

# TensorFlow: 'same' padding available
conv_tf = tf.keras.layers.Conv2D(64, 3, padding='same')

# ONNX handles this conversion automatically, but be aware for manual conversions

Activation Functions

# Some activations have different default parameters

# PyTorch LeakyReLU
leaky_pt = torch.nn.LeakyReLU(negative_slope=0.01)

# TensorFlow LeakyReLU
leaky_tf = tf.keras.layers.LeakyReLU(alpha=0.01)  # Different parameter name

# GELU implementations may differ slightly
gelu_pt = torch.nn.GELU()  # Uses tanh approximation by default
gelu_tf = tf.keras.activations.gelu  # Uses erf (exact) by default

4.6 Advanced Conversion Techniques

Custom Layer Mapping

from onnx_tf.handlers.backend import BackendHandler
from onnx_tf.handlers.handler import onnx_op
import tensorflow as tf

# Define custom ONNX operator handler
@onnx_op("CustomOp")
class CustomOpHandler(BackendHandler):
    @classmethod
    def version_1(cls, node, **kwargs):
        # Implement custom conversion logic
        x = kwargs["tensor_dict"][node.inputs[0]]
        alpha = node.attrs.get("alpha", 1.0)

        # Map to TensorFlow operations
        output = tf.multiply(x, alpha)

        return [output]

# Now ONNX models with CustomOp can be converted to TensorFlow

Quantization-Aware Conversion

import torch
import tensorflow as tf

# PyTorch quantization-aware training
model_pt = models.resnet18(pretrained=True)
model_pt.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_prepared = torch.quantization.prepare_qat(model_pt)

# After QAT training...
model_quantized = torch.quantization.convert(model_prepared)

# Export quantized model to ONNX
torch.onnx.export(
    model_quantized,
    dummy_input,
    "resnet18_quantized.onnx",
    opset_version=13
)

# Convert to TFLite with quantization
converter = tf.lite.TFLiteConverter.from_saved_model("model_tf")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

# Representative dataset for calibration
def representative_dataset():
    for _ in range(100):
        yield [np.random.randn(1, 3, 224, 224).astype(np.float32)]

converter.representative_dataset = representative_dataset
tflite_model = converter.convert()

4.7 Troubleshooting Common Conversion Issues

Issue 1: Unsupported Operators

# Problem: Operator not supported in target framework
# Solution: Replace with supported equivalent or implement custom handler

# Example: Replace unsupported operation
class ModelWithUnsupportedOp(torch.nn.Module):
    def forward(self, x):
        # torch.einsum may not be supported
        # Replace with equivalent matmul operations
        return torch.matmul(x, x.transpose(-2, -1))

# Or use ONNX fallback operators
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH
)

Issue 2: Dynamic Shapes

# Problem: Target framework requires fixed input shapes
# Solution: Export with specific shapes or use dynamic shape support

# Fixed shape export
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "model_fixed.onnx",
    input_names=['input'],
    output_names=['output']
    # No dynamic_axes specified
)

# Dynamic shape export (where supported)
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "model_dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch', 2: 'height', 3: 'width'},
        'output': {0: 'batch'}
    }
)

Issue 3: Numerical Precision

# Problem: Different frameworks have different precision defaults
# Solution: Explicitly set dtypes and validate outputs

import numpy as np

def validate_conversion(model_original, model_converted, test_inputs, tolerance=1e-5):
    """
    Validate conversion maintains numerical accuracy
    """
    for i, test_input in enumerate(test_inputs):
        # Original output
        output_original = model_original(test_input).detach().numpy()

        # Converted output
        output_converted = model_converted(test_input).numpy()

        # Compute differences
        abs_diff = np.abs(output_original - output_converted)
        rel_diff = abs_diff / (np.abs(output_original) + 1e-8)

        max_abs_diff = np.max(abs_diff)
        max_rel_diff = np.max(rel_diff)

        print(f"Test {i+1}:")
        print(f"  Max absolute difference: {max_abs_diff:.2e}")
        print(f"  Max relative difference: {max_rel_diff:.2e}")

        if max_abs_diff > tolerance:
            print(f"  ⚠️ Warning: Exceeds tolerance {tolerance}")
        else:
            print(f"  ✓ Within tolerance")

    return True

# Usage
test_inputs = [torch.randn(1, 3, 224, 224) for _ in range(10)]
validate_conversion(model_pytorch, model_tensorflow, test_inputs)

4.8 Conversion Best Practices

1. Always Validate

• Run identical inputs through both models
• Compare outputs with appropriate tolerances
• Test edge cases and boundary conditions
• Validate on representative dataset samples

2. Use Latest Opsets

# Use recent ONNX opsets for better operator coverage
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=15)  # Good
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=9)   # Outdated

3. Simplify Models Before Conversion

from onnxsim import simplify
import onnx

# Simplify ONNX model before converting to other frameworks
model = onnx.load("model.onnx")
model_simplified, check = simplify(model)

if check:
    onnx.save(model_simplified, "model_simplified.onnx")
    # Now convert simplified model to target framework

4. Document Conversion Steps

# Create a conversion manifest
conversion_info = {
    'source_framework': 'pytorch',
    'source_version': '2.0.1',
    'target_framework': 'tensorflow',
    'target_version': '2.12.0',
    'intermediate_format': 'onnx',
    'onnx_opset': 15,
    'conversion_date': '2025-01-15',
    'validation_passed': True,
    'max_numerical_diff': 1.2e-6,
    'notes': 'Converted ResNet-18 for TF Serving deployment'
}

import json
with open('conversion_manifest.json', 'w') as f:
    json.dump(conversion_info, f, indent=2)

4.9 Automation and CI/CD Integration

Automated Conversion Pipeline

#!/usr/bin/env python3
"""
Automated model conversion pipeline with validation
"""

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

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelConverter:
    def __init__(self, source_path, target_format):
        self.source_path = Path(source_path)
        self.target_format = target_format

    def convert(self):
        logger.info(f"Converting {self.source_path} to {self.target_format}")

        # Load source model
        model = torch.load(self.source_path)
        model.eval()

        # Export to ONNX
        onnx_path = self.source_path.with_suffix('.onnx')
        dummy_input = torch.randn(1, 3, 224, 224)

        torch.onnx.export(
            model, dummy_input, str(onnx_path),
            opset_version=15, export_params=True
        )

        logger.info(f"✓ Exported to ONNX: {onnx_path}")

        # Validate ONNX
        onnx_model = onnx.load(str(onnx_path))
        onnx.checker.check_model(onnx_model)
        logger.info("✓ ONNX model validated")

        # Convert to target format
        if self.target_format == 'tensorflow':
            self._convert_to_tensorflow(onnx_path)
        elif self.target_format == 'tflite':
            self._convert_to_tflite(onnx_path)

        logger.info("✓ Conversion completed")

    def _convert_to_tensorflow(self, onnx_path):
        # Implementation here
        pass

    def _convert_to_tflite(self, onnx_path):
        # Implementation here
        pass

# Usage in CI/CD
if __name__ == "__main__":
    converter = ModelConverter("model.pth", "tensorflow")
    converter.convert()

4.10 Framework-Specific Optimizations

TensorRT Optimization (NVIDIA GPUs)

import tensorrt as trt
import pycuda.driver as cuda

# Convert ONNX to TensorRT for maximum GPU performance
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, fp16_mode=True):
    with trt.Builder(TRT_LOGGER) as builder, \
         builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
         trt.OnnxParser(network, TRT_LOGGER) as parser:

        # Configure builder
        config = builder.create_builder_config()
        config.max_workspace_size = 1 << 30  # 1GB

        if fp16_mode:
            config.set_flag(trt.BuilderFlag.FP16)

        # Parse ONNX
        with open(onnx_path, 'rb') as model:
            parser.parse(model.read())

        # Build engine
        engine = builder.build_engine(network, config)

        # Save engine
        with open(engine_path, 'wb') as f:
            f.write(engine.serialize())

        return engine

# Build TensorRT engine
engine = build_engine("model.onnx", "model.trt", fp16_mode=True)
print("✓ TensorRT engine built")