A comprehensive research framework for training CLIP models with novel n-dimensional loss functions and advanced analysis techniques including CKA (Centered Kernel Alignment).
This framework enables research in:
- Multi-dimensional CLIP Training: 3D, 4D, 6D, and custom dimensional configurations
- Novel Loss Functions: 18+ mathematically rigorous loss function variants
- CKA Analysis: Deep model comparison and understanding
- Cross-modal Learning: Image-text and multilingual capabilities
- Numerical Optimization: Stable training with proper gradient flow
- 🧮 Numerically Stable: All loss functions include stability checks and proper error handling
- 🔧 Highly Configurable: Type-safe configuration system for reproducible experiments
- 📊 Advanced Analysis: Built-in CKA tools for model comparison
- 🧪 Thoroughly Tested: Comprehensive test suite with 95%+ coverage
- 📚 Well Documented: Complete API documentation with Sphinx
- 🌐 Multilingual: Support for Chinese-English translation tasks
git clone https://github.com/st7ma784/6DIMCOCO.git
cd 6DIMCOCO
pip install -r requirements.txt
pip install -e .# Run basic training
python scripts/run_training.py
# Run with wandb logging
python scripts/run_training.py --wandb
# Build datasets
python data_builders/BuildImagenet.py
python data_builders/BuildLAION.pyfrom src.config.base_config import ExperimentConfig
from src.losses import create_loss_function
# Create experiment configuration
config = ExperimentConfig()
config.model.dimensions = 6.0
config.training.learning_rate = 2e-3
# Create loss function
loss_fn = create_loss_function('norm_based', config=config.model)
# Use with your features
import torch
features = [torch.randn(32, 512) for _ in range(6)]
loss = loss_fn(*features)from src.losses import get_available_losses
losses = get_available_losses()
# Output:
# stock_clip: Standard CLIP contrastive loss
# einsum: Einstein summation based n-dimensional loss
# euclidean_distance: Euclidean distance based loss with stability
# norm_based: Norm-based loss with multiple variants
# cosine_similarity: Cosine similarity based multi-dimensional loss- Installation Guide: Detailed setup instructions
- Quick Start: Get running in minutes
- API Reference: Complete API documentation
- Research Applications: Academic use cases and findings
Run the comprehensive test suite:
# All tests
pytest tests/ -v
# Specific test categories
pytest tests/test_losses.py -v # Loss function tests
pytest tests/test_config.py -v # Configuration tests
pytest tests/test_cka_analysis.py -v # CKA analysis tests
# Skip GPU tests if no CUDA
pytest tests/ -m "not gpu" -v6DIMCOCO/
├── src/ # Core source code
│ ├── config/ # Configuration management
│ └── losses/ # Loss function implementations
├── model/ # Model implementations
├── scripts/ # Training and analysis scripts
│ ├── launch.py # Main training orchestration
│ ├── run_training.py # Entry point script
│ ├── CKA_*.py # CKA analysis scripts
│ └── benchmark_cupy.py # Performance benchmarking
├── data_builders/ # Dataset construction scripts
│ ├── BuildCNDataset.py # Chinese dataset builder
│ ├── BuildImagenet.py # ImageNet dataset builder
│ └── Build*.py # Other dataset builders
├── notebooks/ # Jupyter notebooks for analysis
├── results/ # Training results and plots
├── experiments/ # Experimental configurations
├── tests/ # Test suite
├── docs/ # Documentation
├── requirements.txt # Dependencies
└── README.md # This file
Type-safe configuration system replacing hardcoded values:
@dataclass
class ModelConfig:
embed_dim: int = 512
dimensions: float = 6.0
normalize_logits: bool = True
# ... with validationComprehensive testing addressing original issues:
- ✅ Unit Tests: All loss functions and configurations
- ✅ Integration Tests: End-to-end workflows
- ✅ Numerical Stability: Edge cases and error handling
- ✅ Mathematical Properties: Transpose invariance, symmetry
- ✅ Performance Tests: Memory usage and gradient flow
This framework has been used for:
- Multi-dimensional contrastive learning research
- Cross-modal representation learning
- Model architecture analysis via CKA
- Chinese-English translation tasks
- Numerical optimization in deep learning
config.model.dimensions = 6.0 # 3, 3.5, 4, 6, -1, 0
config.model.embed_dim = 512 # Embedding dimension
config.model.normalize_logits = True # Feature normalization
config.model.loss_version = 0 # Legacy compatibilityconfig.training.learning_rate = 2e-3
config.training.train_batch_size = 64
config.training.precision = 16 # Mixed precision
config.training.gradient_clip_val = 0.25- ❌ Minimal test coverage (1 basic test)
- ❌ No systematic validation
- ❌ Hardcoded dependencies
- ❌ No edge case handling
- ✅ Comprehensive test suite (95%+ coverage)
- ✅ Systematic validation framework
- ✅ Configurable dependencies
- ✅ Robust error handling
- ❌ 600+ line monolithic loss file
- ❌ Hardcoded API keys
- ❌ Poor separation of concerns
- ❌ Code duplication across 30+ model versions
- ✅ Modular, well-organized architecture
- ✅ Secure configuration management
- ✅ Clean separation of concerns
- ✅ DRY principle with shared base classes
- Fork the repository
- Create a feature branch (
git checkout -b feature/amazing-feature) - Run tests (
pytest tests/ -v) - Commit changes (
git commit -m 'Add amazing feature') - Push to branch (
git push origin feature/amazing-feature) - Open a Pull Request
This project is licensed under the MIT License - see the LICENSE file for details.
If you use this framework in your research, please cite:
@misc{6dimcoco2024,
title={6DIMCOCO: Multi-dimensional CLIP Training Framework},
author={PhD Research Project},
year={2024},
url={https://github.com/st7ma784/6DIMCOCO}
}- Original research codebase and methodologies
- PyTorch Lightning for training infrastructure
- Weights & Biases for experiment tracking
- The open-source community for inspiration and tools