scripts/run_CXR_CT_train.py

import argparse
import datetime
import os
import yaml

import torch

from transformer_maskgit import CTViT
from transformers import BertTokenizer, BertModel
from radir import RADIR

from RadIR_Trainer_CXR_CT import RadIRTrainer

import json
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--name',required=True, type=str)
    parser.add_argument('--config_file', type=str, help='Path to the config file (YAML format)')
    parser.add_argument("--log_dir", type=str, default=None, help="Path to save logs.")

    parser.add_argument('--checkpoint', type=str)
    parser.add_argument('--resume_training', action='store_true')
    parser.add_argument('--allow_partial_load', action='store_true')
    parser.add_argument('--evaluate_before_train', action='store_true')
    parser.add_argument('--use_triplet_loss', type=float, default=1)
    parser.add_argument('--use_infoNCE_loss', type=float, default=0)
    parser.add_argument('--use_image2image_loss', type=float, default=1)
    parser.add_argument('--use_uncon_triplet_loss', type=float, default=1)
    parser.add_argument('--use_uncon_infoNCE_loss', type=float, default=1)
    parser.add_argument('--freeze_pretrained_encoders', action='store_true')
    parser.add_argument('--open_partial_pretrained_encoders', action='store_true')
    parser.add_argument('--open_pretrained_encoders', action='store_true')
    parser.add_argument('--open_vision_encoder', action='store_true')
    parser.add_argument('--open_fusion_module', action='store_true')
    parser.add_argument('--freeze_fusion_module', action='store_true')
    parser.add_argument('--triplet_loss_margin',type=float, default=0.1)
    parser.add_argument('--positive_distance_threshold',type=float, default=0.25, help='triplet中ij>ik超过此阈值时拉大距离')
    parser.add_argument('--positive_threshold',type=float, default=0.75, help='infoNCE中大于此阈值的都视为false negative & triplet中大于此阈值才可以被作为positive')
    parser.add_argument('--negative_threshold',type=float, default=1, help='triplet中小于此阈值才可以被作为negative')
    parser.add_argument('--anatomy_filter', type=str, default='[["left hemidiaphragm","airspace"]["right rib","thyroid gland"]]')
    parser.add_argument('--dataset_names',type=str, nargs='+', default=['MIMIC-CXR','CT_RATE'])
    parser.add_argument('--modality',type=str, nargs='+',default=['2D', '3D'])
    parser.add_argument('--results_folder', type=str, default='/mnt/data-3/users/zhangzeyu/HMClip/RadIR/logs')

    parser.add_argument('--data_train_jsonl', nargs='+',type=str, required=True)
    parser.add_argument('--data_valid_jsonl', nargs='+',type=str, required=True)
    parser.add_argument('--data_train_csv_dir', nargs='+',type=str, required=True)
    parser.add_argument('--data_valid_csv_dir', nargs='+',type=str, required=True)
    parser.add_argument('--data_train_npy_dir', nargs='+',type=str, required=True)
    parser.add_argument('--data_valid_npy_dir', nargs='+',type=str, required=True)
    parser.add_argument('--local_batch_size', nargs='+',type=int, default=[64,8])
    parser.add_argument('--batch_size', nargs='+',type=int, default=[4,1])
    parser.add_argument('--train_valid_length_',type=int, default=2400)
    # Additional unconstrained dataset parameters
    parser.add_argument('--uncon_batch_size', nargs='+', type=int, help='Batch size for unconstrained datasets')
    parser.add_argument('--uncon_batch_size_valid', nargs='+', type=int, help='Batch size for validation of unconstrained datasets')
    parser.add_argument('--uncon_soft_label', action='store_true', help='Use soft labels for unconstrained datasets')
    # 表示文件夹的路径，并不是实际文件的路径
    parser.add_argument('--uncon_similarity_lookup_table_train', nargs='+', type=str, help='Similarity lookup tables for unconstrained training datasets')
    parser.add_argument('--uncon_similarity_lookup_table_valid', nargs='+', type=str, help='Similarity lookup tables for unconstrained validation datasets')
    parser.add_argument('--uncon_data_train_jsonl', nargs='+', type=str, help='JSONL files for unconstrained training datasets')
    parser.add_argument('--uncon_data_valid_jsonl', nargs='+', type=str, help='JSONL files for unconstrained validation datasets')
    parser.add_argument('--uncon_dataset_names', nargs='+', type=str, help='Names of unconstrained datasets')
    parser.add_argument('--uncon_train_filter', type=str, help='2D list of training filters for unconstrained datasets')
    parser.add_argument('--uncon_valid', nargs='+', type=str, help='List of validation dataset names for unconstrained datasets')
    parser.add_argument('--uncon_modality', nargs='+', type=str, help='Modality types (2D/3D) for unconstrained datasets')
    parser.add_argument('--stage1', action='store_true', help='Whether to run stage 1 training')
    parser.add_argument('--stage2', action='store_true', help='Whether to run stage 2 training')
    parser.add_argument('--modal_embedding', action='store_true', help='Whether to use modal embdding')
    parser.add_argument('--train_no_aug', action='store_true', help='Whether to use modal embdding')
    parser.add_argument('--infonce_temp', action='store_true', help='Whether to use infonce temperature')
    parser.add_argument('--lr', type=float, nargs='+', default=1e-4)
    parser.add_argument('--uncon_num_workers', type=int, default=4)
    parser.add_argument('--con_num_workers', type=int, default=8)
    parser.add_argument('--pin_memory', action='store_true')
    parser.add_argument('--num_train_steps', type=int, nargs='+', default=100000)
    parser.add_argument('--warmup_steps', type=int, nargs='+', default=10000)
    parser.add_argument('--save_results_every', type=int, default=10000)
    parser.add_argument('--save_model_every', type=int, default=10000)
    parser.add_argument('--shuffle_train_samples_every', type=int, default=100000)
    parser.add_argument('--valid_max_samples', type=int, default=2000)  # 2k, allow all valid samples
    parser.add_argument('--train_max_samples', type=int, default=12000) # 24k -> 12k
    config = parser.parse_args()
    config.anatomy_filter = json.loads(config.anatomy_filter)  # convert string to list
    if config.config_file:
        with open(config.config_file, 'r') as f:
            loaded_config_dict = yaml.safe_load(f)
        for key, value in loaded_config_dict.items():
            if hasattr(config, key):
                setattr(config, key, value)
    
    results_folder = f'{config.results_folder}/{config.name}'
    tokenizer = BertTokenizer.from_pretrained('microsoft/BiomedVLP-CXR-BERT-specialized',do_lower_case=True)
    text_encoder = BertModel.from_pretrained("microsoft/BiomedVLP-CXR-BERT-specialized")

    image_encoder = CTViT(
        dim = 512,
        codebook_size = 8192,
        image_size = 480,
        patch_size = 20,
        temporal_patch_size = 10,
        spatial_depth = 8,
        temporal_depth = 6,
        cls_depth = 4,
        dim_head = 32,
        heads = 8
    )


    Rad_IR = RADIR(  
        image_encoder = image_encoder,
        text_encoder = text_encoder,
        tokenizer = tokenizer,
        use_triplet_loss = config.use_triplet_loss,
        use_infoNCE_loss = config.use_infoNCE_loss,
        use_uncon_triplet_loss = config.use_uncon_triplet_loss,
        use_uncon_infoNCE_loss = config.use_uncon_infoNCE_loss,
        use_image2image_loss = config.use_image2image_loss,
        infonce_temp = config.infonce_temp,
        positive_threshold = config.positive_threshold,
        negative_threshold = config.negative_threshold,
        positive_distance_threshold = config.positive_distance_threshold,
        dim_text = 768,
        dim_image = 512,
        dim_latent = 512,
        extra_latent_projection = False,         # whether to use separate projections for text-to-image vs image-to-text comparisons (CLOOB)
        use_mlm=False,
        downsample_image_embeds = False,
        use_all_token_embeds = False
    )
    
    assert sum([
        config.open_pretrained_encoders,
        config.open_partial_pretrained_encoders,
        config.freeze_pretrained_encoders,
        config.open_vision_encoder
    ]) == 1, "Exactly one of open_pretrained_encoders, open_partial_pretrained_encoders, freeze_pretrained_encoders must be True"
    
    
    assert sum([
        config.open_fusion_module,
        config.freeze_fusion_module
    ]) == 1, "Exactly one of open_fusion_module, freeze_fusion_module must be True"
    
    if config.freeze_pretrained_encoders:
        if int(os.environ.get("RANK", 0)) == 0:
            print('** MODEL ** Freeze encoders')
        Rad_IR.freeze_text_encoder()
        Rad_IR.freeze_vision_encoder()
    elif config.open_partial_pretrained_encoders:
        if int(os.environ.get("RANK", 0)) == 0:
            print('** MODEL ** Partially open encoders')
        Rad_IR.open_partial_text_encoder()
        Rad_IR.open_partial_vision_encoder()
    elif config.open_pretrained_encoders:
        if int(os.environ.get("RANK", 0)) == 0:
            print('** MODEL ** Open encoders')
        Rad_IR.open_text_encoder()
        Rad_IR.open_vision_encoder()
    elif config.open_vision_encoder:
        Rad_IR.freeze_text_encoder()
        Rad_IR.open_vision_encoder()
    
    if config.open_fusion_module:
        if int(os.environ.get("RANK", 0)) == 0:
            print('** MODEL ** Open fusion module')
        Rad_IR.open_fusion_module()
    elif config.freeze_fusion_module:
        if int(os.environ.get("RANK", 0)) == 0:
            print('** MODEL ** Freeze fusion module')
        Rad_IR.freeze_fusion_module()
        
    # 输出pid
    print(f"最开始 pid={os.getpid()} ",'rank', torch.distributed.get_rank() if torch.distributed.is_initialized() else 'not distributed')
    if config.stage1 and config.stage2:
        print('同时开始两阶段的训练')
    elif config.stage1:
        print('开始第一阶段的训练')
    elif config.stage2:
        print('开始第二阶段的训练')
    else:
        raise ValueError('At least one of stage1 or stage2 must be True')
    trainer = RadIRTrainer(
        Rad_IR,
        tokenizer = tokenizer,
        data_train_jsonl = config.data_train_jsonl,
        data_valid_jsonl = config.data_valid_jsonl,
        data_train_npy_dir = config.data_train_npy_dir,
        data_valid_npy_dir = config.data_valid_npy_dir,
        data_train_csv_dir = config.data_train_csv_dir,
        data_valid_csv_dir = config.data_valid_csv_dir,
        anatomy_filter = config.anatomy_filter,
        dataset_names = config.dataset_names,
        modality = config.modality,
        positive_threshold = config.positive_threshold,
        negative_threshold = config.negative_threshold,
        batch_size = config.batch_size,
        local_batch_size = config.local_batch_size,
        stage1= config.stage1,
        stage2= config.stage2,
        uncon_batch_size= config.uncon_batch_size,
        uncon_batch_size_valid= config.uncon_batch_size_valid,
        uncon_soft_label = config.uncon_soft_label,
        uncon_similarity_lookup_table_train = config.uncon_similarity_lookup_table_train,
        uncon_similarity_lookup_table_valid = config.uncon_similarity_lookup_table_valid,
        uncon_data_train_jsonl = config.uncon_data_train_jsonl,
        uncon_data_valid_jsonl = config.uncon_data_valid_jsonl,
        uncon_dataset_names = config.uncon_dataset_names,
        uncon_train_filter = json.loads(config.uncon_train_filter) if config.uncon_train_filter else None,
        uncon_valid = config.uncon_valid,
        uncon_modality = config.uncon_modality,
        modal_embedding = config.modal_embedding,
        train_no_aug = config.train_no_aug,
        train_valid_length_ = config.train_valid_length_,
        results_folder = results_folder,
        lr = config.lr,
        num_train_steps = config.num_train_steps,
        warmup_steps = config.warmup_steps,
        current_step = 1,
        save_results_every = config.save_results_every,
        save_model_every = config.save_model_every,
        shuffle_train_samples_every = config.shuffle_train_samples_every,
        con_num_workers = config.con_num_workers,
        uncon_num_workers= config.uncon_num_workers,
        pin_memory = config.pin_memory,
        train_max_samples=config.train_max_samples,
        valid_max_samples=config.valid_max_samples
    )
    
    f_path = os.path.join(config.results_folder, config.name, "log.txt")
    os.makedirs(os.path.dirname(f_path), exist_ok=True)

    if int(os.environ.get("RANK", 0)) == 0:
        with open(f_path, 'a') as f:
            # set exp time
            SHA_TZ = datetime.timezone(datetime.timedelta(hours=8),
                                name='Asia/Shanghai')   
            utc_now = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc)
            beijing_now = utc_now.astimezone(SHA_TZ)    # 北京时间
            exp_time = f'{beijing_now.year}-{beijing_now.month}-{beijing_now.day}-{beijing_now.hour}-{beijing_now.minute}'
            f.write(f'{exp_time} \n Configs :\n')
            configDict = config.__dict__
            for eachArg, value in configDict.items():
                f.write(eachArg + ' : ' + str(value) + '\n')
            f.write('\n')
        print(f'Log at {results_folder}')

    if config.checkpoint is not None:
        if config.resume_training:
            trainer.resume_training(config.checkpoint)
        else:
            trainer.load_checkpoint(config.checkpoint, config.allow_partial_load)
            
    if int(os.environ.get("RANK", 0)) == 0:
        print('** MODEL ** The following parameters are FROZEN:')
        for name, param in Rad_IR.named_parameters():
            if not param.requires_grad:
                print(name)
        print('** MODEL ** The following parameters are TRAINABLE:')
        for name, param in Rad_IR.named_parameters():
            if param.requires_grad:
                print(name)
            
        print(f'Start Training')
    
    trainer.train(config.evaluate_before_train)