Vue d'ensemble

PyTorch 2.5, publié le 17 octobre 2024, ajoute le support CuDNN et la compilation régionale pour optimiser partiellement les modèles.

Fonctionnalités principales

CuDNN et compilation régionale

Le backend CuDNN est mieux intégré et la compilation régionale permet de compiler uniquement certaines parties du modèle avec torch.compile.

python 
import torch

class MonModele(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = torch.nn.Linear(256, 128)
        self.decoder = torch.nn.Linear(128, 256)

    def forward(self, x):
        # Compilation régionale : seul l'encoder est compilé
        x = torch.compile(self.encoder)(x)
        return self.decoder(x)

model = MonModele()
print(model(torch.randn(1, 256)).shape)

Sources