- Create an environment with python3.8.16 and activate. Further steps require this environment to be activated.
conda create -n task2box python=3.8.16
conda activate task2box
- Install cuda 11.7 and pytorch2.0.1:
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia
- Install additional packages with:
pip install -r requirements.txt
The following instructions are to be run inside src/
The training for existing datasets can be run with the following. Note that box_dim is specified to be 2 here (which means 2-dim boxes will be used) and trained on case 1. Ten cases are made available for training (case_num [1-10]). feat_type can be one of [clip_gauss, clip_ave, fim].
# For CUB+iNat
python train_hierarchy.py --dataset cubinat --gt_pairs_fp ../data/hierarchy_cubinat.csv --ckpts_dir ../ckpts_cubinat_80 --box_dim 2 --case_num 1 --feat_type "clip_gauss"
# For ImageNet
python train_hierarchy.py --dataset imagenet --gt_pairs_fp ../data/hierarchy_imagenet.csv --ckpts_dir ../ckpts_imagenet_80 --box_dim 2 --case_num 1 --feat_type "clip_gauss"
Similar to hierarchical datasets, case_num and box_dim can be changed as needed. model_type can be changed into one of [Task2Box, linear, mlp]
python train_taskonomy.py --model_type "Task2Box" --case_num 5 --box_dim 2 --ckpts_dir "../ckpts_taskonomy_80"
The following instructions are to be run inside src/
For existing datasets, the following can be run:
# For CUB+iNat
python eval_hierarchy.py --box_dim 2 --case_num 1 --feat_type "clip_gauss" --model_type "Task2Box" --dataset cubinat --gt_pairs_fp ../data/hierarchy_cubinat.csv --ckpts_dir ../ckpts_cubinat_80 --results_dir "../results_cubinat"
# For ImageNet
python eval_hierarchy.py --box_dim 2 --case_num 1 --feat_type "clip_gauss" --model_type "Task2Box" --dataset imagenet --gt_pairs_fp ../data/hierarchy_imagenet.csv --ckpts_dir ../ckpts_imagenet_80 --results_dir "../results_imagenet"
For novel datasets, the following can be run:
# For CUB+iNat
python eval_hierarchy_novel.py --box_dim 2 --case_num 1 --feat_type "clip_gauss" --model_type "Task2Box" --dataset cubinat --gt_pairs_fp ../data/hierarchy_cubinat.csv --ckpts_dir ../ckpts_cubinat_80 --results_dir "../results_cubinat_novel"
# For ImageNet
python eval_hierarchy_novel.py --box_dim 2 --case_num 1 --feat_type "clip_gauss" --model_type "Task2Box" --dataset imagenet --gt_pairs_fp ../data/hierarchy_imagenet.csv --ckpts_dir ../ckpts_imagenet_80 --results_dir "../results_imagenet_novel" --fim_dir ../embeddings/task2vec_imagenet/ --gauss_fp ../embeddings/imagenet_clip_gauss.pickle
Note that the above scripts produce the metrics for a specific case, feature type, model type, and box dim. To compile all results (for multiple cases) in a directory into a single csv file, the following can be run:
python compile_hierarchy_results.py --results_dir "../results_imagenet_novel" --compiled_results_dir "../compiled_results"
The following will output the compiled results in a single file.
python eval_taskonomy.py
If you found this helpful, please cite our paper:
@InProceedings{Daroya_2024_CVPR,
author = {Daroya, Rangel and Sun, Aaron and Maji, Subhransu},
title = {Task2Box: Box Embeddings for Modeling Asymmetric Task Relationships},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2024},
pages = {28827-28837}
}