Aller au contenu principal

Module 3 — nn.Module : structurer un modèle en couches réutilisables

Le module précédent a appris à manipuler des poids W et b à la main. C'est utile pour comprendre l'autograd, insoutenable au-delà de trois lignes. nn.Module est la brique de PyTorch qui encapsule un sous-réseau : ses paramètres, sa passe avant, ses sous-modules, sa sérialisation. Tous les modèles réels — de la régression logistique au ResNet du module 9 — en descendent.

Deux méthodes suffisent : __init__ et forward​

Un nn.Module est une classe Python qui hérite de torch.nn.Module. On y déclare les sous-modules et les paramètres dans __init__, on décrit le calcul dans forward.

import torch
from torch import nn

class ReseauFashionMNIST(nn.Module):
    def __init__(self, dim_cachee: int = 128, nb_classes: int = 10):
        super().__init__()
        self.aplatir = nn.Flatten()
        self.couche1 = nn.Linear(28 * 28, dim_cachee)
        self.activation = nn.ReLU()
        self.couche2 = nn.Linear(dim_cachee, nb_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.aplatir(x)
        x = self.couche1(x)
        x = self.activation(x)
        return self.couche2(x)

modele = ReseauFashionMNIST()
lot = torch.randn(4, 1, 28, 28)
sortie = modele(lot)                 # appelle forward
print(sortie.shape)                  # torch.Size([4, 10])

Deux règles à ne jamais oublier. Le super().__init__() en tête est obligatoire : c'est lui qui initialise les registres internes qui suivent les paramètres et les sous-modules. Et on appelle modele(lot), pas modele.forward(lot) : la méthode __call__ de nn.Module ajoute la gestion des points d'ancrage (hooks) et l'entrée en mode entraînement ou évaluation, indispensables plus tard.

Sous-modules et enregistrement automatique​

Chaque attribut assigné dans __init__ qui est lui-même un nn.Module ou un nn.Parameter est automatiquement enregistré. C'est ce qui permet à modele.parameters() de rassembler toutes les feuilles à entraîner, sans les énumérer à la main.

for nom, parametre in modele.named_parameters():
    print(nom, tuple(parametre.shape))
# couche1.weight (128, 784)
# couche1.bias (128,)
# couche2.weight (10, 128)
# couche2.bias (10,)
Une liste Python ne suffit pas

Écrire self.couches = [nn.Linear(10, 10), nn.Linear(10, 10)] n'enregistre rien : une liste Python n'est pas un Module, et ses paramètres restent invisibles à parameters(). Le modèle ne sera pas entraîné et vous ne verrez aucune erreur. La forme correcte est self.couches = nn.ModuleList([...]) ou nn.Sequential(*couches).

nn.Sequential contre classe dédiée​

Pour un empilement strict de couches, nn.Sequential évite la classe.

modele_court = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28 * 28, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

C'est concis, lisible, et suffisant pour prototyper. Dès qu'il y a une branche, un tenseur qui saute une couche, un calcul intermédiaire, ou qu'on veut nommer les sous-modules pour les geler individuellement au module 9, la classe dédiée devient nécessaire. Elle expose aussi repr(modele) et l'inspection interactive de façon beaucoup plus lisible.

Un compromis fréquent : la classe expose un bloc = nn.Sequential(...) comme sous-module, et forward orchestre plusieurs blocs. On combine alors la concision de Sequential et l'expressivité de la classe.

state_dict : la représentation portable du modèle​

state_dict() renvoie un dictionnaire ordonné qui contient tous les paramètres et les tampons (moyennes courantes de BatchNorm, par exemple). Sa clé est le chemin qualifié dans l'arbre du modèle, sa valeur est un tenseur.

etat = modele.state_dict()
print(list(etat.keys())[:2])
# ['couche1.weight', 'couche1.bias']

torch.save(etat, "modele_fashion.pt")

# Plus tard, ailleurs
modele_recharge = ReseauFashionMNIST()
modele_recharge.load_state_dict(torch.load("modele_fashion.pt"))

C'est ce format qu'on sauvera systématiquement dans les points de contrôle du module 8, et jamais l'objet Python via torch.save(modele). Sauver l'objet crée une dépendance au chemin exact de la classe : déplacer le fichier .py casse le chargement. Sauver state_dict ne dépend que des noms d'attributs, ce qui reste stable.

L'initialisation des poids influence l'entraînement​

PyTorch initialise les couches avec des valeurs par défaut raisonnables, souvent adaptées à la fonction d'activation courante. Ce défaut est suffisant pour Fashion-MNIST. Il ne l'est plus pour un réseau profond avec des activations exotiques ou une géométrie particulière.

def init_kaiming(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_uniform_(module.weight, nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)

modele.apply(init_kaiming)

modele.apply(fn) parcourt récursivement tous les sous-modules et applique la fonction. Kaiming convient à ReLU, Xavier convient à tanh ; choisir sans réfléchir peut suffire à faire diverger un réseau profond, alors qu'un choix cohérent ne coûte rien et évite plusieurs heures de diagnostic.

Une bonne pratique de nommage

Nommer les couches avec leur rôle (aplatir, couche1, tete) rend state_dict lisible et facilite la mise à jour partielle des poids au module 9. Éviter des noms génériques comme net1, layer2, qui deviennent des étiquettes anonymes six mois plus tard.

Mode entraînement contre mode évaluation​

modele.train() et modele.eval() ne modifient qu'un drapeau : self.training sur chaque sous-module. Ce drapeau change ensuite le comportement de certaines couches — Dropout s'active ou non, BatchNorm utilise le lot courant ou ses moyennes courantes.

modele.train()             # avant la boucle d'entraînement
# ...
modele.eval()              # avant la boucle de validation

Ces deux appels ne calculent rien, ne coûtent rien, mais leur oubli est la source la plus fréquente d'un modèle qui « fonctionne » à l'entraînement et échoue en validation. On y insiste au module 5.

Vérification systématique par une passe avant à blanc​

Avant même de démarrer une boucle d'entraînement, on passe un lot factice pour vérifier que les formes s'enchaînent.

modele = ReseauFashionMNIST()
lot_factice = torch.randn(2, 1, 28, 28)
sortie = modele(lot_factice)
print(sortie.shape)                  # torch.Size([2, 10])

nb_params = sum(p.numel() for p in modele.parameters() if p.requires_grad)
print(f"{nb_params:,} paramètres entraînables")

Cet exercice de trente secondes rattrape la moitié des erreurs d'architecture : une dimension mal calculée, une couche oubliée dans forward, un Flatten manquant. Comptez aussi les paramètres : un modèle inattendument gigantesque signale presque toujours un nn.Linear(dim_H * dim_W, ...) sur des cartes de convolution laissées en 4D par erreur.

En résumé​

  • Deux méthodes suffisent dans un nn.Module : __init__ déclare les sous-modules, forward décrit le calcul ; on appelle modele(x), pas modele.forward(x).
  • Les sous-modules doivent être enregistrés via des attributs directs, nn.ModuleList ou nn.Sequential ; une liste Python nue passe sous le radar de parameters().
  • state_dict est la seule sérialisation à retenir : dictionnaire de tenseurs indexés par le chemin qualifié, indépendant du chemin Python de la classe.
  • L'initialisation par défaut suffit à Fashion-MNIST, mais devient décisive sur les réseaux profonds ; modele.apply(fn) propage un choix cohérent partout.

Le module suivant s'occupe des données : comment Dataset et DataLoader alimentent la passe avant en lots correctement mélangés, transformés et découpés entre entraînement, validation et test.