[TMI Submission] UniChest

This is the official repository of UniChest (TMI Submission) with the source code and the pre-trained model weights.

Pre-training

Conquer Stage

CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 train.py --main_ratio 1 --bias_ratio 0 --moe_ratio 0 --output_dir --aws_output_dir

Divide Stage

CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 train.py --main_ratio 0.5 --bias_ratio 0.5 --moe_ratio 1 --output_dir --aws_output_dir --finetune

Arguments

• --output_dir is the directory to save logs
• --aws_output_dir is the directory to save checkpoints
• --finetune is the path of the checkpoint of the Conquer Stage

Pre-trained model weights

The pre-trained model can be downloaded from google drive.

Testing

python test.py --main_ratio 0.5 --bias_ratio 0.5 --aws_output_dir --test_data --save_result_dir

Arguments

• --aws_output_dir is the path of the checkpoint
• --test_data is dataset name
• --save_result_dir is the path to save ground truth and prediction results