# GRAM/HyperGRAM Paper Experiments

This directory contains the code and experimental configurations for the GRAM/HyperGRAM paper.

## Directory Structure

```
paper_clean/
├── model/               # Complete model implementations (GRAM, HyperGRAM, encoders)
├── evaluation/          # Evaluation metrics and pipelines
├── utils/              # Utility functions (volume computation, logging, etc.)
├── data/               # Data loading and preprocessing code
├── scripts/
│   ├── stage1/         # Stage 1 pretraining scripts (VAST150K)
│   │   ├── fixed_weights/     # Fixed weight experiments (4 scripts)
│   │   ├── curvature/         # Curvature experiments (8 scripts)
│   │   ├── validation_task/   # Validation task experiments (2 scripts)
│   │   └── baseline/          # Baseline experiment (1 script)
│   └── stage2/         # Stage 2 fine-tuning scripts
│       └── msrvtt/            # MSR-VTT retrieval experiments (1 script)
├── configs/
│   ├── pretrain/              # Stage 1 pretrain configs (15 files)
│   └── finetune/              # Stage 2 finetune configs (1 file)
├── run.py                     # Main training entry point
├── pyproject.toml             # Python dependencies
└── README.md                  # This file
```

## Experiments

### 1. Fixed Weights Experiments (4 experiments)

Testing different weight combinations between Euclidean and Hyperbolic representations:
- **Fixed 0.1**: 10% Euclidean, 90% Hyperbolic
- **Fixed 0.3**: 30% Euclidean, 70% Hyperbolic
- **Fixed 0.7**: 70% Euclidean, 30% Hyperbolic
- **Fixed 0.9**: 90% Euclidean, 10% Hyperbolic

**Scripts**: `scripts/stage1/fixed_weights/train_stage1_pretrain_fixed_{0.1,0.3,0.7,0.9}.sh`

**Configs**: `configs/pretrain/pretrain_hybrid_fixed_{0.1,0.3,0.7,0.9}_vast150k.json`

### 2. Curvature Experiments (8 experiments)

Testing different fixed curvature values in hyperbolic space: 0.05, 0.3, 0.5, 0.7, 1.2, 1.5, 2.0, 9.0

Lower values (0.05) approximate Euclidean space, higher values (9.0) represent strongly curved hyperbolic space.

**Scripts**: `scripts/stage1/curvature/train_stage1_pretrain_curve_{value}.sh`

**Configs**: `configs/pretrain/pretrain_hybrid_curve_{value}_vast150k.json`

### 3. Validation Task Experiments (2 experiments)

Comparing validation tasks during pretraining:
- **TV**: Text-Video retrieval only
- **TVA**: Text-Video-Audio retrieval

**Scripts**: `scripts/stage1/validation_task/train_stage1_pretrain_curve_1.0_val_{tv,tva}.sh`

**Configs**: `configs/pretrain/pretrain_hybrid_curve_1.0_val_{tv,tva}_vast150k.json`

### 4. Baseline (1 experiment)

Baseline with VATEX validation set.

**Script**: `scripts/stage1/baseline/train_stage1_pretrain_vatex_val.sh`

**Config**: `configs/pretrain/pretrain_hybrid_vast150k_vatex_val_paper.json`

## Two-Stage Training Pipeline

### Stage 1: Pretraining on VAST150K

Dataset: VAST150K (150K video-text-audio triplets)

Training: 4 GPUs, batch size 32/GPU, 15 epochs, lr=1e-4

```bash
cd /path/to/paper_clean
bash scripts/stage1/curvature/train_stage1_pretrain_curve_1.0.sh
```

### Stage 2: Fine-tuning on MSR-VTT

Dataset: MSR-VTT video-text retrieval

Training: 8 GPUs, batch size 32/GPU, 10 epochs, lr=5e-5

```bash
cd /path/to/paper_clean
export PRETRAIN_DIR=/path/to/stage1/checkpoint
bash scripts/stage2/msrvtt/train_stage2_finetune_msrvtt_pre.sh
```

## Running Experiments

### Prerequisites

Environment:
```bash
uv sync
```

Data:
- VAST150K: `~/data/VAST150K/`
- MSR-VTT: `~/data/msrvtt/`
- Pretrained weights: `../pretrained_weights/`

Scripts support SLURM job submission with `/scratch/` for faster I/O.

### Running a Stage 1 Experiment

```bash
# Example: Run curvature=1.0 experiment
cd /path/to/paper_clean
sbatch scripts/stage1/curvature/train_stage1_pretrain_curve_1.0.sh
```

The script will:
1. Request 4 GPUs with 40GB memory each
2. Copy datasets to scratch space for faster I/O
3. Run distributed training with PyTorch DDP
4. Save checkpoints to `output2/pretrain_hybrid_curve_1.0_vast150k/`
5. Log training progress to `train.log`

### Running a Stage 2 Experiment

```bash
# Set the pretrained checkpoint directory
export PRETRAIN_DIR=/path/to/output2/pretrain_hybrid_curve_1.0_vast150k

# Submit fine-tuning job
cd /path/to/paper_clean
sbatch scripts/stage2/msrvtt/train_stage2_finetune_msrvtt_pre.sh
```

## Hardware Requirements

**Minimum Requirements**:
- 4x NVIDIA A100 40GB GPUs (for stage 1)
- 8x NVIDIA A100 40GB GPUs (for stage 2)
- 256GB RAM
- 1TB disk space (for datasets and checkpoints)

**Recommended Setup**:
- SLURM cluster with NVMe scratch storage
- InfiniBand networking for multi-node training
- Fast shared filesystem for checkpoint storage

## Expected Training Times

**Stage 1 Pretraining** (VAST150K, 15 epochs):
- 4x A100 40GB: ~24-30 hours per experiment
- Total for all 15 experiments: ~450 GPU-hours

**Stage 2 Fine-tuning** (MSR-VTT, 10 epochs):
- 8x A100 40GB: ~3-4 hours per experiment
- Total: ~30 GPU-hours

## Key Configuration Parameters

### Hybrid Geometry Parameters

- `euclidean_weight`: Weight for Euclidean component (0.0-1.0)
- `hyperbolic_weight`: Weight for Hyperbolic component (0.0-1.0)
- `curvature`: Fixed curvature value for Poincaré ball (>0)
- `learnable_curvature`: Whether to learn curvature (true/false)
- `learnable_weights`: Whether to learn geometry weights (true/false)

### Training Parameters

- `train_batch_size`: Batch size per GPU for training
- `val_batch_size`: Batch size per GPU for validation
- `learning_rate`: Initial learning rate
- `num_train_epochs`: Number of training epochs
- `warmup_ratio`: Warmup ratio for learning rate scheduler

### Validation Parameters

- `validation_task`: Task for validation ("tv", "tva", or "all")
- `val_datasets`: Datasets used for validation (e.g., "vatex", "msrvtt")

## Output Structure

Each experiment creates an output directory with:

```
output2/pretrain_hybrid_{experiment_name}_vast150k/
├── checkpoint_best.pt      # Best checkpoint based on validation
├── checkpoint_last.pt      # Last checkpoint
├── train.log              # Training logs
└── config.json            # Training configuration
```

## Notes

- All scripts use `uv run python` for execution
- Distributed training uses `torch.distributed.launch`
- Mixed precision (FP16) is enabled by default
- Wandb logging is automatically started for rank 0 process
- Log files are saved as `train.log` in the output directory

## Support

For questions or issues, check the configuration files in `configs/` for parameter details.
