🤖 machine learning

Towards Migrating Neural Network Implementations

This paper proposes and validates an automated approach for migrating neural network code between deep learning frameworks, such as PyTorch and TensorFlow, by utilizing a pivot model to create an abstraction that ensures functional equivalence while overcoming the challenges of manual migration.

Original authors: Nadia Daoudi, Ivan Alfonso, Jordi Cabot

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nadia Daoudi, Ivan Alfonso, Jordi Cabot

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have built a magnificent, complex machine (a Neural Network) using a specific set of blueprints and tools from one country (let's say, TensorFlow). Now, your company decides to move to a new country where everyone uses a different set of tools and speaks a slightly different language (PyTorch).

The problem? You can't just pick up your machine and move it. The gears don't fit, the instructions are in the wrong language, and the parts are shaped differently. Usually, you'd have to tear the whole machine apart and rebuild it from scratch by hand, which takes forever and is prone to errors.

This paper presents a robotic translator that does this rebuilding for you automatically.

Here is how their "magic translator" works, broken down into simple steps:

1. The "Universal Translator" (The Pivot Model)

Instead of trying to translate directly from Language A to Language B (which is messy because they are so different), the authors invented a Universal Language (called the BESSER-NN Metamodel).

Think of this like a universal adapter for electrical plugs.

Step 1: They take your original machine (TensorFlow code) and break it down into its basic parts (layers, connections, settings). They translate these parts into the "Universal Language."
Step 2: They take that Universal Language and translate it into the new country's language (PyTorch code).

This ensures that no matter where you start or where you end up, the core "soul" of the machine remains the same.

2. The Three-Step Process

The authors describe their method as a three-stage assembly line:

Stage 1: The X-Ray (AST Extraction): They use a tool to take an "X-ray" of the original code. This turns the messy text code into a clear, organized tree diagram (called an Abstract Syntax Tree) that shows exactly how the machine is built.
Stage 2: The Translation (Transformation): They look at that tree diagram and rewrite it into their "Universal Language" blueprint. This is where they handle the tricky parts, like making sure the new machine knows how big the input data is, even if the old code didn't explicitly say so.
Stage 3: The Construction (Code Generation): Finally, they use a template system (like a "fill-in-the-blanks" form) to write the new code in the target language (PyTorch or TensorFlow) based on that universal blueprint.

3. The Tricky Hurdles (Challenges)

The paper admits that this wasn't easy. They had to solve three specific puzzles:

The "Missing Piece" Puzzle: In the old language (TensorFlow), some parts of the machine automatically figure out their size based on what you feed them. In the new language (PyTorch), you must write down the size explicitly. The authors' robot had to do some detective work to "guess" and fill in those missing sizes so the new machine wouldn't crash.
The "Hidden Instruction" Puzzle: Sometimes, the activation function (a switch that turns the machine on) is hidden inside a variable name rather than written out clearly. The robot had to be smart enough to find that hidden name and replace it with the correct, explicit instruction for the new language.
The "Left-Handed vs. Right-Handed" Puzzle: Imagine one country packs boxes with the handle on the left, and the other packs them with the handle on the right. If you just move the box, it won't fit on the shelf. TensorFlow and PyTorch organize data differently (one puts the "color channels" at the end, the other at the start). The authors added a "re-arranging" step to flip the data around so the new machine can read it correctly without changing the actual picture.

4. Did It Work? (The Results)

The authors tested their robot on five famous neural network designs (like AlexNet and VGG16). They tried moving them back and forth between TensorFlow and PyTorch.

Did it build? Yes. The new code ran without errors.
Did it think the same? Yes. They fed the exact same pictures and numbers into both the old machine and the new machine. The results were almost identical (differing by less than one-billionth of a percent).
Did it learn the same? Yes. When they trained both machines on real-world data (like recognizing cats vs. dogs or movie reviews), they achieved nearly the same accuracy scores.

The Bottom Line

The paper claims they have built a tool that can automatically translate neural network code from one major framework to another. This saves developers from having to manually rewrite thousands of lines of code, ensuring the new version works exactly like the old one, just speaking a different "language."

They also mention that this fits into a larger system where software is designed using models, but their main focus is strictly on the code translation itself, not on how the AI is used in hospitals or other specific real-world applications.

Technical Summary: Towards Migrating Neural Network Implementations

Problem Statement

The rapid advancement of neural networks (NNs) has led to a proliferation of deep learning (DL) frameworks (e.g., PyTorch, TensorFlow). While organizations adopt these frameworks to leverage specific features, performance characteristics, or community support, they often face the need to migrate existing NN implementations to a different framework later in the lifecycle. This migration is driven by factors such as declining performance, evolving requirements, or the deprecation of support for specific frameworks (e.g., Hugging Face deprecating TensorFlow support).

Currently, there is a significant lack of automated approaches tailored specifically for migrating NN code. Existing literature primarily focuses on converting trained, ready-to-use models (artifacts) rather than the source code itself. Consequently, organizations are forced to manually rewrite NN code from scratch, a process that is time-consuming, error-prone, and hinders the flexibility to adopt new features while reusing existing logic.

Methodology

The authors propose a Model-Driven Engineering (MDE) approach to automatically migrate NN code between DL frameworks. The core of the method relies on a pivot model to abstract the source code before generating the target code. The process consists of three main steps:

1. AST Extraction

The process begins by parsing the source framework's code (Python) to extract its Abstract Syntax Tree (AST). The authors utilize Python's ast library to represent the code as a tree of nodes, where nodes correspond to NN components such as layers, tensor operations, and hyperparameters. This step supports both Sequential and Subclassing architectures found in PyTorch and TensorFlow.

2. Transformation (Model-to-Model)

The extracted AST is transformed into a framework-agnostic representation using the BESSER-NN metamodel. This metamodel, part of the BESSER low-code platform, serves as the pivot model.

Mapping: Nodes from the source AST (e.g., a TensorFlow Dense layer) are mapped to their equivalents in the BESSER-NN model (e.g., a Linear layer).
Attributes: Hyperparameters, training configurations, and dataset definitions are also mapped to the pivot model.
Handling Architectural Differences: The approach employs four distinct transformation modules to handle the four combinations of source/target architectures (PyTorch/TensorFlow $\times$ Subclassing/Sequential).
Symbolic Tracing: A critical aspect of this step involves symbolically tracing the source code to infer attributes not explicitly defined in the AST (e.g., input shapes in TensorFlow), ensuring the pivot model contains all necessary information for the target framework.

3. Code Generation (Model-to-Text)

Once the BESSER-NN pivot model is populated, code generators produce the final implementation in the target framework.

Preprocessing: Python scripts process the BESSER-NN model to resolve cross-references and dependencies between layers.
Template Engine: The Jinja template engine is used to map the abstract concepts in the pivot model to the specific syntax of the target framework (PyTorch or TensorFlow).
Architecture Support: Four code generators are developed to support all architecture combinations, ensuring that complex designs (like defining a network within another class) are preserved.

Key Contributions and Technical Challenges

The paper addresses several specific technical challenges inherent to cross-framework migration:

Missing Attributes (Input Shapes): In TensorFlow, input shapes are often inferred automatically, whereas PyTorch requires explicit definition. The authors solved this by adding a symbolic tracing step during the transformation phase to infer and populate missing input shape attributes (e.g., in_features, in_channels) in the pivot model before generating PyTorch code.
Dynamic Activation Functions: TensorFlow allows activation functions to be assigned dynamically to variables and passed to layers, whereas PyTorch typically expects explicit layer definitions. The migration approach handles this by verifying the activation function value; if it is a known string, a standard layer is generated. If it is dynamic, the generator creates a custom function call that resolves the activation at runtime, preserving the original logic.
Channel Ordering Conventions: TensorFlow uses "channel-last" ordering, while PyTorch uses "channel-first." To maintain functional equivalence without altering the input data format, the approach inserts permute tensor operations. These operations reorder tensor dimensions before and after convolutional layers in the target code to align with the framework's native convention while keeping the data flow consistent with the source.
System Integration: The paper demonstrates how the BESSER-NN metamodel can be integrated with general software modeling languages like UML. This allows NNs to be defined as behavior implementations within broader system models (e.g., as methods in a state machine), facilitating the migration of NNs as part of a larger smart system migration.

Experimental Evaluation

The authors evaluated the approach using five distinct neural networks from the literature: AlexNet, VGG16, LSTM, CNN-RNN, and TF-Tutorial. These networks cover 12 distinct layer types and a mix of Sequential and Subclassing architectures.

The evaluation addressed two Research Questions (RQs):

RQ1 (Migration Success): The approach successfully migrated all five networks across all supported scenarios (PyTorch $\leftrightarrow$ TensorFlow, Sequential $\leftrightarrow$ Subclassing). The generated code executed without structural errors when fed dummy input data.
RQ2 (Functional Preservation): The authors verified that the migrated networks produce outputs functionally equivalent to the originals.
- Random Inputs: Using identical weights and random inputs, the Maximum Absolute Difference (MAD) between source and target outputs was negligible (up to $5.96 \times 10^{-8}$ ).
- Benchmark Datasets: Networks were trained on CIFAR-10, SVHN, IMDB, and SST-2. The prediction accuracy of the migrated networks was highly consistent with the source networks, with the largest observed difference being 0.022 (for VGG16 on CIFAR-10).

Significance and Claims

The paper claims that its primary significance lies in shifting the focus from migrating trained models to migrating source code. Unlike existing tools (e.g., ONNX, MMdnn, CoreML) that convert static, trained artifacts, this approach migrates the definition of the network itself.

This distinction is significant because it:

Enables iterative experimentation, allowing developers to inspect, modify, and retrain the model within the new framework.
Supports model-driven pipelines, allowing NNs to be treated as first-class citizens within broader software engineering workflows.
Avoids the "black box" nature of artifact conversion, preserving the ability to leverage new features and best practices of the target framework.

The authors conclude that their approach effectively bridges the gap between DL frameworks, reducing the manual effort required for modernization while maintaining functional equivalence. They note that future work will involve extending the metamodel to support more advanced AI concepts (e.g., Transformers, Graph Neural Networks) and developing refactoring transformations to improve NN structure.