PyTorch Model Zoo 🔥

Production-quality implementations of modern deep learning architectures.

Models

Model	Paper	Task	Status
ViT	Dosovitskiy et al. 2020	Image Classification	✅
GPT-2	Radford et al. 2019	Language Modeling	✅
DDPM	Ho et al. 2020	Image Generation	✅
ResNet	He et al. 2015	Image Classification	✅
BERT	Devlin et al. 2019	NLU	🚧

Highlights

• Clean, readable implementations
• Proper weight initialization
• Mixed precision training (AMP)
• Gradient accumulation support
• Checkpoint saving/loading
• TensorBoard logging

Quick Start

pip install -r requirements.txt

# Train ViT on CIFAR-10
python train.py --model vit --dataset cifar10 --epochs 100

# Train GPT-2 on custom text
python train.py --model gpt2 --data data/shakespeare.txt --epochs 50

# Generate with trained GPT-2
python generate.py --checkpoint checkpoints/gpt2_best.pt --prompt "The meaning of life" --max_tokens 200

Project Structure

models/
├── vit.py              # Vision Transformer
├── gpt2.py             # GPT-2 language model
├── ddpm.py             # Denoising Diffusion Probabilistic Model
├── resnet.py           # ResNet family
└── components/
    ├── attention.py    # Multi-head attention
    ├── embeddings.py   # Position/token embeddings
    └── mlp.py          # Feed-forward blocks
train.py                # Universal training script
generate.py             # Inference / generation
config.py               # Model configurations
requirements.txt