Minimal Implementation of Scalable Rectified Flow Transformers

Ce dépôt contient une implémentation minimale des modèles de flux rectifiés. J’ai adopté l’approche de l’entraînement avec l’architecture SD3. Cette fois, j’ai décidé de diviser le fichier en 2 parties : l’implémentation du modèle et le code réel, mais vous n’avez pas besoin de regarder le code du modèle. Tout est toujours autonome, minimal et facile à pirater. Rien de compliqué ne se passe si vous comprenez les mathématiques.

1. Simple Rectified Flow, pour les débutants

Installez torch, pil, torchvision

Exécutez

python rf.py

pour entraîner le modèle sur MNIST à partir de zéro. Si vous êtes cool et que vous voulez entraîner CIFAR à la place, vous pouvez le faire. Sur le 63e epoch, votre sortie devrait ressembler à quelque chose comme ça.

2. Massive Rectified Flow, muP Support

Ceci est pour les gigachads qui veulent entraîner Imagenet à la place. Ne vous inquiétez pas ! IMO Imagenet est le nouveau MNIST, et nous utiliserons mon jeu de données pour cela. Tout d’abord, allez dans le dossier avancé, téléchargez le jeu de données.

cd advancedpip install hf_transfer # juste installez ceci.bash download.sh

Cela ne devrait pas prendre plus de 5 minutes si votre réseau est correct. Exécutez

bash run.sh

pour entraîner le modèle. Cela entraînera Imagenet à partir de zéro, fera une recherche en grille muP pour trouver le bassin aligné pour la fonction de perte, vous déverrouillerez le transfert LR sans tir pour les modèles de flux rectifiés !

Cela utilise plusieurs techniques et bases de code que j’ai développées au fil des ans. C’est un mélange naturel de , ,

Citations

Si vous utilisez ce matériel, veuillez citer ce dépôt avec les informations suivantes :

@misc{ryu2024minrf,  author={Simo Ryu},  title={minRF: Minimal Implementation of Scalable Rectified Flow Transformers},  year=2024,  publisher={Github},  url={https://github.com/cloneofsimo/minRF},}

Source de l’article