PyTorch → TFLite Convert

This post was migrated from Tistory. You can find the original here.

An alley in Itaewon

PyTorch to TFLite Github

https://github.com/cornpip/pt_to_tflite

PyTorch to TFLite Docker Image

https://hub.docker.com/r/cornpip77/tf_213_converter

This is an image built on top of tensorflow:2.13.0-gpu, with the following packages installed:

pip install torch torchvision onnx onnx-tf tensorflow-addons tensorflow-probability

Example

As a conversion example, let’s convert mobilenet_v2 to TFLite.

Just follow the MobileNet Test Script in the README.

// torch model (mobilenet_v2-b0353104.pth)
258: 0.8303 (Samoyed)
259: 0.0699 (Pomeranian)
261: 0.0130 (keeshond)
231: 0.0108 (collie)
257: 0.0099 (Great Pyrenees)

// tflite model (mobilenet_v2.tflite)
258: 0.8692 (Samoyed)
259: 0.0498 (Pomeranian)
261: 0.0243 (keeshond)
257: 0.0163 (Great Pyrenees)
231: 0.0047 (collie)

// tflite no optimizer model (no quantization)
258: 0.8595 (Samoyed)
259: 0.0568 (Pomeranian)
261: 0.0252 (keeshond)
257: 0.0177 (Great Pyrenees)
231: 0.0060 (collie)

After converting following the MobileNet Test Script and checking the results, you can see there’s a difference in the inference output.

Even without quantization, there’s a difference between the torch model’s inference results and the converted model’s.

In my experience, the behavior varies depending on the trained model.

Some models see their GT accuracy drop to 60-70% after quantization,

while others show almost no drop in GT accuracy. Of course, even in that case, there’s still a difference in the resulting probability values.

Before using the converted TFLite on another platform

Before using a converted tflite model on another platform, it’s recommended to first verify the TFLite inference within a Python environment.

If there’s a large gap in GT accuracy at this point, you should assume something is wrong and start debugging. (Check GT accuracy first, then look at differences in probability values.)

• build_model isn’t identical, or
• the model is heavily affected by quantization, or
• something was missed during the conversion process

The reason it’s a good idea to check this in Python first is that,

even if you implement the same preprocessing pipeline, there are inevitably subtle differences between platforms—resize interpolation method, float↔int casting policy, how the original image is loaded, etc.
and these differences affect inference.

So if you don’t first verify in Python that the TFLite conversion was done correctly, debugging becomes much harder later on.

Walking through the conversion code

main

The conversion flow is as follows.

PyTorch → ONNX → TensorFlow → TFLite

def main() -> None:
    args = parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[device] Using {device}")

    # 1) Prepare model + load weights
    model = build_model(
        model_type=args.model,
        num_classes=args.num_classes,
        backbone=args.backbone,
        use_pretrained_backbone=args.use_pretrained_backbone,
    )
    state_dict = load_state_dict_flexible(args.pt_path, device=torch.device("cpu"))
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    print("[state_dict] missing:", missing)
    print("[state_dict] unexpected:", unexpected)
    model.eval().to(device)

    # 2) PyTorch → ONNX
    dummy_input = torch.randn(1, 3, args.input_height, args.input_width, device=device)
    onnx_model_path = f"./onnx/{args.result_name}.onnx"
    export_onnx(model, dummy_input, onnx_model_path)
    print(f"[onnx] saved to {onnx_model_path}")

    # 3) ONNX → TensorFlow SavedModel
    saved_model_dir = f"./saved_model/{args.result_name}"
    print("[tf] converting onnx → saved_model ...")
    onnx_to_tf_nhwc(onnx_model_path, saved_model_dir)
    print(f"[tf] saved_model ready at {saved_model_dir}")

    # 4) TensorFlow → TFLite
    tflite_model_path = f"./tflite/{args.result_name}.tflite"
    print("[tflite] converting saved_model → tflite ...")
    tf_to_tflite(saved_model_dir, tflite_model_path, optimize=True)
    print(f"[tflite] saved to {tflite_model_path}")
    
    
 def export_onnx(model: nn.Module, dummy_input: torch.Tensor, onnx_path: str) -> None:
    """PyTorch → ONNX"""
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        opset_version=11,
    )
    
  def tf_to_tflite(saved_model_dir: str, tflite_path: str, optimize: bool = True) -> None:
    """TensorFlow SavedModel → TFLite"""
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
    if optimize:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
    tflite_model = converter.convert()
    with open(tflite_path, "wb") as f:
        f.write(tflite_model)

• If you don’t want quantization, you can remove that part of tf_to_tflite() (tf.lite.Optimize.DEFAULT).

build_model

The most important thing is to accurately reproduce the exact same layers as the trained model.

import torch.nn as nn
from torchvision import models
from torchvision.models import EfficientNet_B4_Weights, ResNet50_Weights

def build_model(model_type: str, num_classes: int, backbone: str, use_pretrained_backbone: bool) -> nn.Module:
    if model_type == "resnet50":
        model = models.resnet50(weights=ResNet50_Weights.DEFAULT)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
        return model

    if model_type == "efficientnet_b4":
        model = models.efficientnet_b4(weights=EfficientNet_B4_Weights.DEFAULT)
        model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        return model

    if model_type == "custom_head":
        return CustomHeadModel(
            backbone_name=backbone,
            num_classes=num_classes,
            pretrained=use_pretrained_backbone,
        )

    raise ValueError(f"Unsupported model type: {model_type}")
    
class CustomHeadModel(nn.Module):
    """
    Example custom model.
    """
    
    def __init__(self, backbone_name: str = "efficientnet_b4", num_classes: int = 3, pretrained: bool = False):
        super().__init__()

        if "efficientnet" in backbone_name:
            base = getattr(models, backbone_name)(
                weights=models.EfficientNet_B4_Weights.DEFAULT if pretrained else None
            )
            in_features = base.classifier[1].in_features
            base.classifier = nn.Identity()
        elif "resnet" in backbone_name:
            base = getattr(models, backbone_name)(
                weights=models.ResNet50_Weights.DEFAULT if pretrained else None
            )
            in_features = base.fc.in_features
            base.fc = nn.Identity()
        else:
            raise ValueError(f"Unsupported backbone: {backbone_name}")

        self.backbone = base
        self.bn = nn.BatchNorm1d(in_features)
        self.head = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        feat = self.backbone(x)
        if feat.dim() == 4:
            feat = torch.flatten(feat, 1)
        feat = self.bn(feat)
        return self.head(feat)

• An example of a model where only the final fc layer is set (resnet50, efficientnet_b4)
• An example of a model with a few more layers stacked in front of the fc layer (custom_head)

onnx_to_tf_nhwc

TFLite input is NHWC. We modify the interface to match that.

import tensorflow as tf
import onnx
from onnx_tf.backend import prepare

def onnx_to_tf_nhwc(onnx_model_path, saved_model_dir):
    # Load ONNX model
    onnx_model = onnx.load(onnx_model_path)
    tf_rep = prepare(onnx_model)
    tf_rep.export_graph(saved_model_dir)

    model = tf.saved_model.load(saved_model_dir)
    concrete_func = model.signatures["serving_default"]
    input_tensor = concrete_func.inputs[0]
    input_shape = input_tensor.shape.as_list()  # [1, C, H, W]
    nhwc_shape = [input_shape[0], input_shape[2], input_shape[3], input_shape[1]]

    @tf.function(input_signature=[tf.TensorSpec(shape=nhwc_shape, dtype=tf.float32)])
    def new_serving_fn(inputs):
        nchw_input = tf.transpose(inputs, [0, 3, 1, 2])
        outputs = concrete_func(nchw_input)
        return outputs

    tf.saved_model.save(model, saved_model_dir, signatures={'serving_default': new_serving_fn})

• After converting ONNX (NCHW) to a TF SavedModel, we wrap it so that only the input is received as NHWC, and save it again.
• The internal operations stay in NCHW; only the external interface is changed to NHWC, so the channel order doesn’t get mixed up during TFLite conversion/inference.