pytorch-patterns

Modèles de Développement PyTorch

Modèles idiomatiques PyTorch et bonnes pratiques pour construire des applications d'apprentissage profond robustes, efficaces et reproductibles.

Quand Activer

• Écriture de nouveaux modèles PyTorch ou de scripts d’entraînement
• Révision de code d’apprentissage profond
• Débogage de boucles d’entraînement ou de pipelines de données
• Optimisation de l’utilisation de la mémoire GPU ou de la vitesse d’entraînement
• Mise en place d’expériences reproductibles

Principes Fondamentaux

1. Code Indépendant du Périphérique

Toujours écrire du code fonctionnant à la fois sur CPU et GPU sans coder en dur le périphérique.

# Bon : Indépendant du périphérique
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

# Mauvais : Périphérique codé en dur
model = MyModel().cuda()  # Plante si pas de GPU
data = data.cuda()

2. Reproductibilité Avant Tout

Définir toutes les graines aléatoires pour des résultats reproductibles.

# Bon : Configuration complète de reproductibilité
def set_seed(seed: int = 42) -> None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# Mauvais : Pas de contrôle de graine
model = MyModel()  # Poids différents à chaque exécution

3. Gestion Explicite des Formes

Toujours documenter et vérifier les formes des tenseurs.

# Bon : Passage avant annoté avec les formes
def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (batch_size, canaux, hauteur, largeur)
    x = self.conv1(x)    # -> (batch_size, 32, H, L)
    x = self.pool(x)     # -> (batch_size, 32, H//2, L//2)
    x = x.view(x.size(0), -1)  # -> (batch_size, 32*H//2*L//2)
    return self.fc(x)    # -> (batch_size, num_classes)

# Mauvais : Pas de suivi des formes
def forward(self, x):
    x = self.conv1(x)
    x = self.pool(x)
    x = x.view(x.size(0), -1)  # Quelle est la taille ?
    return self.fc(x)           # Est-ce que ça marchera ?

Modèles d'Architecture

Structure Propre de nn.Module

# Bon : Module bien organisé
class ImageClassifier(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64 * 16 * 16, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# Mauvais : Tout dans forward
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = F.conv2d(x, weight=self.make_weight())  # Crée le poids à chaque appel !
        return x

Initialisation Correcte des Poids

# Bon : Initialisation explicite
def _init_weights(self, module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.BatchNorm2d):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model = MyModel()
model.apply(model._init_weights)

Modèles de Boucle d'Entraînement

Boucle d'Entraînement Standard

# Bon : Boucle d’entraînement complète avec bonnes pratiques
def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    scaler: torch.amp.GradScaler | None = None,
) -> float:
    model.train()  # Toujours passer en mode entraînement
    total_loss = 0.0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad(set_to_none=True)  # Plus efficace que zero_grad()

        # Entraînement en précision mixte
        with torch.amp.autocast("cuda", enabled=scaler is not None):
            output = model(data)
            loss = criterion(output, target)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

Boucle de Validation

# Bon : Évaluation correcte
@torch.no_grad()  # Plus efficace que d'encapsuler dans un bloc torch.no_grad()
def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()  # Toujours passer en mode évaluation — désactive dropout, utilise les stats BN cumulées
    total_loss = 0.0
    correct = 0
    total = 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        total_loss += criterion(output, target).item()
        correct += (output.argmax(1) == target).sum().item()
        total += target.size(0)

    return total_loss / len(dataloader), correct / total

Modèles de Pipeline de Données

Dataset Personnalisé

# Bon : Dataset propre avec indications de type
class ImageDataset(Dataset):
    def __init__(
        self,
        image_dir: str,
        labels: dict[str, int],
        transform: transforms.Compose | None = None,
    ) -> None:
        self.image_paths = list(Path(image_dir).glob("*.jpg"))
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        img = Image.open(self.image_paths[idx]).convert("RGB")
        label = self.labels[self.image_paths[idx].stem]

        if self.transform:
            img = self.transform(img)

        return img, label

Configuration Efficace de DataLoader

# Bon : DataLoader optimisé
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # Mélanger pour l'entraînement
    num_workers=4,           # Chargement parallèle
    pin_memory=True,         # Transfert CPU->GPU plus rapide
    persistent_workers=True, # Garder les workers actifs entre les époques
    drop_last=True,          # Tailles de batch constantes pour BatchNorm
)

# Mauvais : Valeurs par défaut lentes
dataloader = DataLoader(dataset, batch_size=32)  # num_workers=0, pas de pin_memory

Collate Personnalisé pour Données de Longueur Variable

# Bon : Remplir les séquences dans collate_fn
def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
    sequences, labels = zip(*batch)
    # Remplir à la longueur max du batch
    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels)

dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

Modèles de Point de Contrôle (Checkpoint)

Sauvegarde et Chargement de Points de Contrôle

# Bon : Point de contrôle complet avec tout l'état d'entraînement
def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: str,
) -> None:
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, path)

def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint

# Mauvais : Sauvegarde uniquement des poids du modèle (impossible de reprendre l’entraînement)
torch.save(model.state_dict(), "model.pt")

Optimisation des Performances

Entraînement en Précision Mixte

# Bon : AMP avec GradScaler
scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
    with torch.amp.autocast("cuda"):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

Checkpointing de Gradient pour Grands Modèles

# Bon : Échanger du calcul contre de la mémoire
from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Recalculer les activations pendant la rétropropagation pour économiser la mémoire
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

torch.compile pour la Vitesse

# Bon : Compiler le modèle pour une exécution plus rapide (PyTorch 2.0+)
model = MyModel().to(device)
model = torch.compile(model, mode="reduce-overhead")

# Modes : "default" (sûr), "reduce-overhead" (plus rapide), "max-autotune" (le plus rapide)

Référence Rapide : Idiomes PyTorch

Idiome	Description
model.train() / model.eval()	Toujours définir le mode avant entraînement/évaluation
torch.no_grad()	Désactiver les gradients pour l'inférence
optimizer.zero_grad(set_to_none=True)	Effacement plus efficace des gradients
.to(device)	Placement tenseur/modèle indépendant du périphérique
torch.amp.autocast	Précision mixte pour une vitesse ×2
pin_memory=True	Transfert de données CPU→GPU plus rapide
torch.compile	Compilation JIT pour la vitesse (2.0+)
weights_only=True	Chargement sécurisé du modèle
torch.manual_seed	Expériences reproductibles
gradient_checkpointing	Échanger du calcul contre de la mémoire

Anti-Modèles à Éviter

# Mauvais : Oublier model.eval() pendant la validation
model.train()
with torch.no_grad():
    output = model(val_data)  # Dropout toujours actif ! BatchNorm utilise les stats du batch !

# Bon : Toujours définir le mode évaluation
model.eval()
with torch.no_grad():
    output = model(val_data)

# Mauvais : Opérations sur place qui cassent l'autograd
x = F.relu(x, inplace=True)  # Peut casser le calcul des gradients
x += residual                  # L'ajout sur place casse le graphe autograd

# Bon : Opérations sans modification sur place
x = F.relu(x)
x = x + residual

# Mauvais : Déplacer les données sur GPU à plusieurs reprises dans la boucle d’entraînement
for data, target in dataloader:
    model = model.cuda()  # Déplace le modèle À CHAQUE itération !

# Bon : Déplacer le modèle une seule fois avant la boucle
model = model.to(device)
for data, target in dataloader:
    data, target = data.to(device), target.to(device)

# Mauvais : Utiliser .item() avant backward
loss = criterion(output, target).item()  # Détache du graphe !
loss.backward()  # Erreur : impossible de rétropropager avec .item()

# Bon : Appeler .item() uniquement pour les logs
loss = criterion(output, target)
loss.backward()
print(f"Perte : {loss.item():.4f}")  # .item() après backward est correct

# Mauvais : Ne pas utiliser torch.save correctement
torch.save(model, "model.pt")  # Sauvegarde tout le modèle (fragile, non portable)

# Bon : Sauvegarder state_dict
torch.save(model.state_dict(), "model.pt")

Rappel : Le code PyTorch doit être indépendant du périphérique, reproductible et économe en mémoire. En cas de doute, profiler avec torch.profiler et vérifier la mémoire GPU avec torch.cuda.memory_summary().