# UniDense: Unleashing Diffusion Models with Meta-Routers for Universal Few-Shot Dense Prediction

## Requirements
* Python 3.9
* [CUDA 11.1](https://developer.nvidia.com/cuda-toolkit)
* [PyTorch 1.10.2](https://pytorch.org)
* torchvision
* timm

## Dataset Preparation
Download Taskonomy Dataset (tiny split) from the official github page [here](https://github.com/StanfordVL/taskonomy/tree/master/data). You may download data of depth_euclidean, depth_zbuffer, edge_occlusion, keypoints2d, keypoints3d, normal, principal_curvature, reshading, segment_semantic, and rgb. 

## Usage
### Multi-Task Pre-Training for the MoE Backbone
```bash
python -m torch.distributed.launch --nproc_per_node=8 --use_env train_mtl.py \
    --data_dir '/xxx/data/taskonomy_tiny' \
    --output_dir '/xxx/output/UniDense/work_dirs' \
    --log_dir '/xxx/output/UniDense/log_dir' \
    --pathdir_model '/xxx/model' \
    --task_fold 1 \
    --batch_size 2 \
    --epochs 50 \
    --blr 1e-4 \
    --weight_decay 0.05 \
    --warmup_epochs 5 \
    --model 'UniDense' \
    --drop_path 0.1 \
    --exp-name 'UniDense_MTL' \
    --num_workers 4 \
    --replica_factor 3
```

### Meta-Training for Meta-Routers
Specify 'path_ckpt_pretrain' with the model obtained from the multi-task pre-training stage.
```bash
python -m torch.distributed.launch --nproc_per_node=8 --use_env train_meta.py \
    --data_dir '/xxx/data/taskonomy_tiny' \
    --output_dir '/xxx/output/UniDense/work_dirs' \
    --log_dir '/xxx/output/UniDense/log_dir' \
    --pathdir_model '/xxx/model' \
    --path_ckpt_pretrain '/xxx/checkpoint_xxx.pth' \
    --task_fold 1 \
    --batch_size 2 \
    --episodes 3000 \
    --blr 5e-4 \
    --weight_decay 0.05 \
    --model 'UniDense' \
    --drop_path 0.1 \
    --exp-name 'UniDense_meta' \
    --num_workers 4 \
    --replica_factor 3
```

### Fine-Tuning on Novel Few-Shot Tasks
Specify 'path_ckpt_pretrain' with the model obtained from the meta-training stage.
```bash
python -m torch.distributed.launch --nproc_per_node=8 --use_env train_ft.py \
    --data_dir '/xxx/data/taskonomy_muleshoe' \
    --output_dir '/xxx/output/UniDense/work_dirs' \
    --log_dir '/xxx/output/UniDense/log_dir' \
    --pathdir_model '/xxx/model' \
    --path_ckpt_pretrain '/xxx/checkpoint_xxx.pth' \
    --task_fold 1 \
    --batch_size 2 \
    --steps_FT 30 \
    --test_episodes 1000 \
    --blr 5e-4 \
    --weight_decay 0.05 \
    --model 'UniDense' \
    --drop_path 0.1 \
    --exp-name 'UniDense_test' \
    --num_workers 4 \
    --replica_factor 3
```
