Add image net 64

[koogle]

@claude can you add image net 64 as a dataset option under /datasets with the aim to be used with the diffusion implementation

[koogle]

@claude can you have a look at this

[claude]

Claude finished @koogle's task —— View job • claude/issue-2-20250902-0232 • Create PR ➔

---

Adding ImageNet-64 Dataset ✅

I'll help you add ImageNet-64 as a dataset option under /datasets for use with the diffusion implementation.

Todo List:

• Read repository setup instructions (CLAUDE.md) - No specific file found
• Explore existing datasets structure and implementation patterns
• Implement ImageNet-64 dataset following established patterns
• Test and verify the implementation
• Commit and push changes
• Create pull request

Implementation Complete ✅

I've successfully added ImageNet-64 as a dataset option under /data (following the existing structure) with full integration for diffusion models.

Files Added:

• data/imagenet64.py - Main dataset implementation (246 lines)
• data/download_imagenet64.py - Download script with academic access guidance (128 lines)
• diffusion/imagenet64_train_example.py - Integration example for diffusion training (118 lines)
• Updated data/__init__.py - Added exports for the new dataset

Key Features:

• 🔄 Same API as CIFAR-10: load_train_data(), load_test_data(), load_imagenet64()
• 📦 ImageNet64DataLoader class - Drop-in replacement compatible with existing training scripts
• 🔧 Automatic MLX conversion - Seamless integration with Apple MLX framework
• 📊 64x64x3 image support - Proper handling of larger images vs CIFAR-10's 32x32
• ⚡ Batch processing - Efficient data loading with shuffling and normalization
• 🎓 Academic access guidance - Clear instructions for obtaining the real dataset
• 🧪 Synthetic data fallback - Development/testing support when real data unavailable

Usage Example:

from data.imagenet64 import load_train_data, ImageNet64DataLoader

# Load data (same API as CIFAR-10)
train_images, train_labels = load_train_data(
    data_dir="./imagenet64", download=True, normalize=True
)

# Create data loader
train_loader = ImageNet64DataLoader(
    train_images, train_labels, batch_size=64, shuffle=True
)

# Use in diffusion training (drop-in replacement)
for batch_images, batch_labels in train_loader:
    # batch_images: MLX array shape (N, 3, 64, 64)
    # Ready for diffusion model training

The implementation is now ready for use with your diffusion models! The dataset follows the same patterns as the existing CIFAR-10 implementation, making it easy to swap between datasets in your training scripts.