Cette deuxième partie du tutoriel sur le modèle DiT (Diffusion Transformer) approfondit l’application des concepts introduits précédemment, cette fois-ci sur un jeu de données plus complexe : EMNIST. L’objectif reste inchangé — transformer un bruit gaussien (distribution normale) en images reconnaissables — mais le défi est désormais de générer des chiffres et des lettres en niveaux de gris, représentés par des tenseurs de forme (28, 28), bien plus complexes que les simples points en 2D manipulés auparavant.
Le jeu de données EMNIST, téléchargé via PyTorch, contient des paires étiquette-image où les étiquettes correspondent à des caractères (lettres ou chiffres). La version "balanced" utilisée ici exclut les classes redondantes, comme les majuscules et minuscules identiques (par exemple, le "O" et le "o"), pour éviter les doublons. Cependant, les images brutes posent un problème visuel : elles apparaissent inversées et tournées de 90 degrés. Pour les corriger, une série de transformations est appliquée via `torchvision.transforms`, incluant une rotation fixe de -90 degrés, un retournement horizontal systématique, et une conversion en tenseur. Ces opérations, bien que conçues pour l’augmentation de données, permettent ici de normaliser l’orientation des caractères.
Un obstacle technique survient lors de l’affichage : les tenseurs générés par `ToTensor()` adoptent une forme (1, 28, 28) — avec un canal unique pour le noir et blanc — incompatible avec `matplotlib.imshow`, qui attend (28, 28) ou (28, 28, 1). La solution consiste à réorganiser les dimensions du tenseur avec `permute(1, 2, 0)`, rendant les images exploitables. Une fois ces ajustements effectués, les échantillons affichent correctement des lettres et chiffres lisibles, prêts pour l’entraînement.
La phase d’apprentissage réutilise une architecture similaire à la première partie, adaptée aux nouvelles dimensions des données. Le modèle est entraîné sur 600 époques avec des lots de 256 images, en optimisant une fonction de perte basée sur la vitesse de diffusion (la différence entre le bruit ajouté et la prédiction du modèle). À chaque itération, un bruit est progressivement mélangé aux images originales selon un paramètre temporel aléatoire, et le modèle apprend à inverser ce processus. Les poids sont mis à jour via l’optimiseur AdamW, avec un taux d’apprentissage fixé à 6e-4. Bien que le code ne détaille pas l’architecture précise du modèle DiT, il souligne l’importance de la gestion des dimensions et des transformations pour traiter des données image, tout en conservant la logique de diffusion introduite précédemment.