Portage de modèle : Migrer de PyTorch à JAX/Flax

Vous avez passé des jours à ajuster les hyperparamètres de votre modèle sur PyTorch, consommant des heures de calcul. Mais voilà : pour la mise en production, vous devez exploiter la puissance brute des TPU Google, ce qui exige d'utiliser JAX/Flax. Faut-il tout ré-entraîner de zéro ? Heureusement, non. C'est ce qu'on appelle le portage de modèle, et voici comment je m'y suis pris pour le réaliser sans perdre une miette de performance.

Qu'est-ce que le portage de modèle et pourquoi est-ce possible ?

Pour comprendre le portage, il faut se rappeler ce qu'est vraiment un réseau de neurones une fois l'entraînement terminé. Ce n'est plus un processus dynamique : c'est une immense suite d'opérations mathématiques figées, principalement des multiplications et des additions de matrices.

Toute "l'intelligence" acquise par votre modèle est stockée sous forme de simples nombres (des float32 ou float16) : les poids (weights) et les biais (biases). PyTorch, JAX, TensorFlow... ne sont finalement que des interfaces qui ordonnent à une puce d'exécuter ces opérations.

Porter un modèle, c'est donc simplement extraire ces nombres pour les injecter dans un modèle équivalent dans un autre framework. C'est possible précisément parce que les mathématiques sous-jacentes sont universelles. D'ailleurs, cette universalité dépasse largement les frameworks classiques. Puisqu'un modèle n'est qu'un fichier contenant des matrices de nombres, le portage peut tout à fait se faire en dehors de tout framework existant. Si vous avez besoin d'aller en bas niveau pour des raisons de performance extrême, ou que vous souhaitez changer complètement de langage de programmation (par exemple passer de Python à Rust, Zig ou C++), il vous suffit de lire ce fichier de poids et de recoder vous-même les boucles de multiplications matricielles. Le modèle n'a pas besoin de PyTorch ou de JAX pour "exister".

Bien sûr, pour des cas standards de transfert entre écosystèmes, des formats intermédiaires comme ONNX existent pour automatiser une partie du travail, mais comprendre le portage manuel reste la clé pour garder un contrôle total et ne pas se retrouver bloqué face à une architecture un peu particulière. La seule vraie difficulté du portage, c'est de comprendre comment chaque framework décide de ranger ces nombres dans la mémoire.

Pour mettre ces principes en pratique et voir concrètement comment les choses se passent, j'ai décidé de m'attaquer à un petit modèle CNN.

L'architecture de notre CNN

Pour rendre les choses concrètes, regardons le modèle que je vais porter. Il s'agit d'un petit réseau de neurones convolutif (SmallCNN) conçu pour classifier des vêtements du jeu de données Fashion-MNIST (des images en niveaux de gris de 28x28 pixels réparties en 10 catégories). L'architecture est classique : deux blocs Convolution + Pooling pour extraire les features visuelles, suivis d'un aplatissement et de deux couches denses (Linéaires) pour prendre la décision finale.

Observer la structure des modèles

Avant de toucher au code de portage, il est indispensable de visualiser clairement la structure de notre modèle et la forme des tenseurs qui y transitent. Sous PyTorch, j'utilise torchsummary dans ce but :

from torchsummary import summary
# Permet de visualiser la structure : couches, tailles d'entrée/sortie, et forme des poids
summary(model, input_size=(64, 1, 28, 28), col_names=["input_size", "output_size", "num_params", "kernel_size"])

De l'autre côté, je crée le même modèle sous Flax. Pour visualiser son architecture avec le format NHWC attendu, on utilise nnx.tabulate :

import jax.numpy as jnp
from flax import nnx
# tabulate permet de mieux visualiser l'architecture sous JAX
print(nnx.tabulate(flax_model, jnp.ones((1, 28, 28, 1))))

Mais si, en plus de la structure, j'inspecte les poids bruts de ce modèle Flax vierge :

flax_state = nnx.state(flax_model)
for key, value in flax_state.items():
    print(f"{key}: {value}")

On voit des poids initialisés aléatoirement, mais surtout, tous les biais sont exactement à zéro. C'est l'état par défaut d'un modèle qui n'a rien appris. Notre but va être de remplacer ces valeurs par celles de PyTorch.

Pour bien comprendre ce qui va se passer lors du portage, voici un tableau qui compare directement comment les données se transforment étape par étape d'un côté sous PyTorch, et de l'autre sous Flax. Vous y verrez immédiatement que si les dimensions sont identiques, leur ordre change totalement à cause du format NCHW pour PyTorch et NHWC pour Flax :

Ce tableau illustre parfaitement le problème central : bien que les tailles de lot (64) et les dimensions spatiales (28x28, puis 14x14, puis 7x7) restent les mêmes, la position des canaux (mis en gras) bouge constamment. Et comme si cela ne suffisait pas, regardons l'étape Flatten : les deux frameworks aboutissent au même chiffre (3136), mais ils n'y arrivent pas en lisant la grille 7x7x64 dans le même ordre ! C'est ce détail qui va nous causer le plus de soucis.

Le problème pratique : NCHW contre NHWC

Comme le montre le tableau, la différence fondamentale de stockage des images oppose le format NCHW de PyTorch (Batch, Channels, Hauteur, Largeur) au format NHWC de JAX/Flax (Batch, Hauteur, Largeur, Channels).

Maintenant que l'on voit clairement ce que l'on a des deux côtés, comment fait-on le pont entre les deux ?

Les trois transformations mathématiques du portage

Pour copier les poids sans corrompre le modèle, notre moteur de portage doit appliquer trois transformations très précises.

1. Les couches Denses (Linear)

C'est l'étape la plus directe. La seule différence réside dans l'ordre de la matrice de poids. PyTorch stocke sa matrice sous la forme (Sortie, Entrée), tandis que Flax l'exige en (Entrée, Sortie), ce qui se résout par une simple transposée mathématique pour inverser les lignes et les colonnes. Il faut également penser à mapper les noms lors du transfert, car PyTorch appelle ce poids weight là où Flax l'appelle kernel.

2. Les Convolutions

Puisque le format global de l'image change de NCHW vers NHWC, les filtres de la convolution doivent logiquement suivre le même mouvement. PyTorch stocke le filtre en (Canaux_Sortie, Canaux_Entrée, Hauteur, Largeur), mais Flax exige de mettre les dimensions spatiales en premier : (Hauteur, Largeur, Canaux_Entrée, Canaux_Sortie). On corrige cela en déplaçant les axes du tenseur avec la fonction permute(2, 3, 1, 0), en n'oubliant pas là encore de renommer le weight en kernel.

3. Le piège de l'aplatissement (Quand la Convolution rencontre la couche Dense)

C'est l'étape où 90% des portages échouent silencieusement, produisant un modèle qui donne des résultats aléatoires. Après nos couches de convolution, on a un tenseur 3D (par exemple : 64 canaux, 7x7 pixels). Pour l'envoyer à la couche Dense, il faut l'aplatir en une ligne unique de 3136 éléments. Mais PyTorch et JAX ne parcourent pas ce tenseur dans le même ordre.

Prenons un exemple visuel avec un minuscule tenseur de 2x2 pixels ayant 2 canaux ($C_1$ et $C_2$). En mémoire, ce bloc 3D ressemble à ça :

$$ \text{Image} = \begin{bmatrix} [C_1, C_2]_{P_1} & [C_1, C_2]_{P_2} \\ [C_1, C_2]_{P_3} & [C_1, C_2]_{P_4} \end{bmatrix} $$

Sous PyTorch (Channel-major) : Le canal est placé au début (NCHW). L'algorithme va faire varier la position spatiale le plus rapidement, et le canal le plus lentement. Il va donc lire tout le canal 1, puis tout le canal 2. Le vecteur aplati est :

$$ V_{\text{PyTorch}} = \begin{bmatrix} C_1P_1 \\ \mathbf{C_1P_2} \\ C_1P_3 \\ C_1P_4 \\ \mathbf{C_2P_1} \\ C_2P_2 \\ C_2P_3 \\ C_2P_4 \end{bmatrix} $$

Sous JAX (Spatial-major) : Le canal est à la fin (NHWC). L'algorithme va faire varier le canal le plus rapidement, et la position spatiale le plus lentement. Il lit tous les canaux du pixel 1, puis tous les canaux du pixel 2. Le vecteur aplati devient :

$$ V_{\text{JAX}} = \begin{bmatrix} C_1P_1 \\ \mathbf{C_2P_1} \\ C_1P_2 \\ C_2P_2 \\ C_1P_3 \\ C_2P_3 \\ C_1P_4 \\ C_2P_4 \end{bmatrix} $$

Regardez bien les valeurs en gras à la deuxième ligne. Dans PyTorch, la deuxième case représente le canal 1 du pixel 2 ($C_1P_2$). Dans JAX, la deuxième case représente le canal 2 du pixel 1 ($C_2P_1$). Si vous vous contentez de transposer la matrice de la couche Dense pour passer de PyTorch à Flax, la ligne numéro 2 de cette matrice de poids (qui a appris à traiter $C_1P_2$ sous PyTorch) va se retrouver connectée à $C_2P_1$ sous JAX. Les connexions sont totalement perverties.

Notre moteur doit donc faire plus qu'une simple transposée : il doit calculer une permutation précise pour réordonner les lignes de la couche Dense, afin que la ligne entraînée pour $C_1P_2$ vienne se placer exactement à la 3ème position (là où JAX attend $C_1P_2$).

Pour automatiser tout ça, j'ai construit un véritable moteur de traduction.

Plongée dans le moteur de portage

Pour construire ce moteur, je me suis d'ailleurs fortement inspiré des excellents exemples de la documentation JAXStack, en l'adaptant pour le rendre complètement récursif et générique. J'ai codé une classe Torch2Flax où les deux premières étapes se traduisent par ces fonctions d'aide :

def _conv_params_permute(name, torch_param):
    if name == "weight":
        # (out_c, in_c, H, W) -> (H, W, in_c, out_c)
        return torch_param.permute(2, 3, 1, 0).contiguous()
    return torch_param

def _linear_params_permute(name, torch_param):
    if name == "weight":
        # (out_feat, in_feat) -> (in_feat, out_feat)
        return torch_param.permute(1, 0).contiguous()
    return torch_param

Le cœur du moteur utilise un dictionnaire pour faire le lien entre les objets PyTorch et Flax, en appliquant la bonne fonction de permutation et en gérant le renommage weight -> kernel :

class Torch2Flax:
    def __init__(self):
        self._last_conv_out_channels = None

    modules_mapping_info = {
        nn.Conv2d: {
            "type": nnx.Conv,
            "params_mapping": {"weight": "kernel", "bias": "bias"},
            "params_transform": _conv_params_permute,
        },
        nn.Linear: {
            "type": nnx.Linear,
            "params_mapping": {"weight": "kernel", "bias": "bias"},
            "params_transform": _linear_params_permute,
        },
        # ... (mapping aussi pour Identity, Flatten, ReLU qui n'ont pas de poids)
    }

Comment ce moteur sait-il quelles couches comparer ? J'ai construit une marche arborée récursive. La méthode copy_module descend dans l'arborescence de PyTorch. Si elle voit un conteneur (comme nn.Sequential), elle avance enfant par enfant. Si elle voit un module avec des attributs (comme self.conv1), elle cherche l'attribut du même nom dans le modèle Flax et y plonge. Quand elle tombe enfin sur une couche de base (Conv ou Linear), elle appelle la copie de paramètres.

C'est dans cette copie qu'on gère le piège de l'aplatissement. Le code détecte si une couche Linéaire fait suite à une Convolution, puis reconstruit mathématiquement l'ordre spatial de JAX :

# Au sein de la méthode de copie des paramètres
if (isinstance(torch_nn_module, nn.Linear)
        and torch_key == "weight"
        and self._last_conv_out_channels is not None):
    C = self._last_conv_out_channels
    in_features = torch_nn_module.in_features
    if in_features > C and in_features % C == 0:
        remaining = in_features // C
        H = int(math.isqrt(remaining))
        W = H
        if H * W == remaining:
            perm = torch.empty(in_features, dtype=torch.long)
            # On recrée l'ordre exact pour passer de Channel-major à Spatial-major
            for c_idx in range(C):
                for h_idx in range(H):
                    for w_idx in range(W):
                        cm_idx = c_idx * H * W + h_idx * W + w_idx # Index PyTorch
                        sm_idx = (h_idx * W + w_idx) * C + c_idx    # Index JAX
                        perm[sm_idx] = cm_idx
            torch_value = torch_value[perm] # Réordonnancement de la matrice !

# On convertit enfin le tenseur PyTorch modifié en tableau JAX natif, puis on l'injecte directement pour écraser les anciens poids
nnx_param.set_value(jnp.asarray(torch_value.data))

Vérification et test de vérité

Après avoir exécuté le script de portage (t2f.copy_module(model, flax_model)), je relance ma commande d'inspection des poids Flax :

flax_state = nnx.state(flax_model)
for key, value in flax_state.items():
    print(f"{key}: {value}")

Les zéros ont disparu. Les biais ont maintenant des valeurs non nulles, et les poids correspondent parfaitement à ceux de PyTorch (aux permutations d'axes près). L'intelligence a réellement migré.

Il ne reste plus qu'à tester le modèle avec de vraies données, sans oublier de transposer les images du jeu de test de NCHW vers NHWC à l'entrée :

for images, labels in test_loader:
    images = images.numpy().transpose(0, 2, 3, 1) # NCHW -> NHWC
    outputs = flax_model(jnp.array(images))
    predictions = jnp.argmax(outputs, axis=1)
    # ... calcul de l'accuracy

Résultat : 91.61% en PyTorch, et 91.61% en Flax. Le portage est un succès total.

Conclusion

En définitive, porter un modèle d'un framework à l'autre n'est pas de la magie noire, mais de la logique pure. Dès lors qu'on saisit qu'un réseau de neurones n'est qu'un empilement de matrices, le portage devient un simple exercice de traduction spatiale. Cette compétence vous rend totalement agnostique : vous n'êtes plus lié à un seul outil, et vous pouvez même exporter vos poids vers des langages comme Rust ou Zig pour un contrôle bas niveau absolu.

Toutefois, il faut rester prudent : ici, nous avons manipulé un petit et simple CNN. Pour des modèles beaucoup plus imposants, comme des architectures de type Transformer (LLM, Vision Transformers), la traduction peut devenir nettement plus ardue, où la gestion des erreurs d'arrondi, des normalisations spécifiques et des couches d'attention se transforme en un véritable casse-tête ! Mais comprendre ces mécanismes de base reste la meilleure des clés.

Liens Utiles

Pour approfondir vos connaissances et explorer des outils avancés, voici quelques ressources :

• Mon code
• Porting a PyTorch model to JAX de JAXStack
• Convert PyTorch models to Flax
• Tips on converting pytorch models to flax