""" Puzzle Tuning Script ver： Feb 11th 14:00 Paper: https://arxiv.org/abs/2311.06712 Code: https://github.com/sagizty/PuzzleTuning Ref: MAE https://github.com/facebookresearch/mae Step 1: PreTraining on the ImagetNet-1k style dataset (others) Step 2: Domain Prompt Tuning (PuzzleTuning) on Pathological Images (in ImageFolder) Step 3: FineTuning on the Downstream Tasks This is the training code for step 2 Pre-training Experiments: DP (data-parallel bash) python PuzzleTuning.py --batch_size 64 --blr 1.5e-4 --epochs 200 --accum_iter 2 --print_freq 2000 --check_point_gap 50 --input_size 224 --warmup_epochs 20 --pin_mem --num_workers 32 --strategy loop --PromptTuning Deep --basic_state_dict /data/saved_models/ViT_b16_224_Imagenet.pth --data_path /root/datasets/All DDP (distributed data-parallel bash) for one machine with 12 GPU python -m torch.distributed.launch --nproc_per_node=12 --nnodes 1 --node_rank 0 PuzzleTuning.py --DDP_distributed --batch_size 64 --blr 1.5e-4 --epochs 200 --accum_iter 2 --print_freq 2000 --check_point_gap 50 --input_size 224 --warmup_epochs 20 --pin_mem --num_workers 32 --strategy loop --PromptTuning Deep --basic_state_dict /data/saved_models/ViT_b16_224_Imagenet.pth --data_path /root/datasets/All update: Use "--seg_decoder" parameter to introduce segmentation networks swin_unet for Swin-Unet """ import argparse import datetime import json import numpy as np import os import time from pathlib import Path import torch import torch.backends.cudnn as cudnn from tensorboardX import SummaryWriter import torchvision.transforms as transforms import torchvision.datasets as datasets import timm # assert timm.__version__ == "0.3.2" # version check import timm.optim.optim_factory as optim_factory import SSL_structures.misc as misc from SSL_structures.misc import NativeScalerWithGradNormCount as NativeScaler from utils.schedulers import patch_scheduler, ratio_scheduler from SSL_structures import models_mae, SAE from SSL_structures.engine_pretrain import train_one_epoch def main(args): # choose encoder for timm basic_encoder = args.model[4:] # choose decoder version args.model = args.model + '_decoder' if args.seg_decoder is not None else args.model # note decoder args.model_idx = args.model_idx + args.model + '_' + args.seg_decoder if args.seg_decoder is not None \ else args.model_idx + args.model # note PromptTuning args.model_idx = args.model_idx + '_Prompt_' + args.PromptTuning + '_tokennum_' + str(args.Prompt_Token_num) \ if args.PromptTuning is not None else args.model_idx # fix the seed for reproducibility if args.DDP_distributed: misc.init_distributed_mode(args) seed = args.seed + misc.get_rank() else: seed = args.seed torch.manual_seed(seed) np.random.seed(seed) # set GPUs cudnn.benchmark = True device = torch.device(args.device) # cuda # simple augmentation transform_train = transforms.Compose([ transforms.RandomResizedCrop(args.input_size, scale=(0.8, 1.0), interpolation=3, ratio=(1. / 1., 1. / 1.)), # 3 is bicubic # transforms.Resize(args.input_size), transforms.RandomVerticalFlip(), transforms.RandomHorizontalFlip(), transforms.ToTensor(), # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) dataset_train = datasets.ImageFolder(os.path.join(args.data_path), transform=transform_train) # , 'train' print('dataset_train:', dataset_train) # Train data if args.DDP_distributed: # args.DDP_distributed is True we use distributed data parallel(DDP) num_tasks = misc.get_world_size() # use misc to set up DDP global_rank = misc.get_rank() # get the rank of the current running sampler_train = torch.utils.data.DistributedSampler( dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True) print("Sampler_train = %s" % str(sampler_train)) enable_DistributedSampler = True batch_size_for_Dataloader = args.batch_size else: # Data parallel(DP) instead of distributed data parallel(DDP) global_rank = 0 sampler_train = torch.utils.data.RandomSampler(dataset_train) enable_DistributedSampler = False batch_size_for_Dataloader = args.batch_size * torch.cuda.device_count() # set log on the main process if global_rank == 0 and args.log_dir is not None: args.log_dir = os.path.join(args.log_dir, args.model_idx) os.makedirs(args.log_dir, exist_ok=True) log_writer = SummaryWriter(log_dir=args.log_dir) # Tensorboard print('Task: ' + args.model_idx) print("Use", torch.cuda.device_count(), "GPUs!") print('job AImageFolderDir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) print("{}".format(args).replace(', ', ',\n')) else: log_writer = None # output_dir if args.output_dir is not None: args.output_dir = os.path.join(args.output_dir, args.model_idx) os.makedirs(args.output_dir, exist_ok=True) os.makedirs(os.path.join(args.output_dir, 'figs'), exist_ok=True) print('Training output files will be at', args.output_dir) else: print('no out put path specified!') raise data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, # the shuffle=True is already set in the sampler batch_size=batch_size_for_Dataloader, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True) # define the model if args.model[0:3] == 'mae': if args.basic_state_dict is not None: # Transfer-learning try: if args.basic_state_dict == 'timm': basic_model = timm.create_model('vit_base_patch' + str(16) + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() print('MAE Transfer-learning with timm') else: basic_state_dict = torch.load(args.basic_state_dict) if 'model' in basic_state_dict: basic_state_dict = basic_state_dict['model'] except: print('erro in args.basic_state_dict:', args.basic_state_dict) if args.PromptTuning is not None: print( 'In PromptTuning, the basic_state_dict is required, without specification now, timm loaded.\n') # timm model name basic_encoder basic_model = timm.create_model(basic_encoder + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() else: basic_state_dict = None print('MAE Restart with a empty backbone') else: print('MAE Transfer-learning with:', args.basic_state_dict) else: if args.PromptTuning is not None: print('In PromptTuning, the basic_state_dict is required, without specification now, timm loaded.\n') # timm model name basic_encoder basic_model = timm.create_model(basic_encoder + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() else: basic_state_dict = None print('MAE Restart with a empty backbone') # mae-vit-base-patch16 model = models_mae.__dict__[args.model](img_size=args.input_size, norm_pix_loss=args.norm_pix_loss, prompt_mode=args.PromptTuning, Prompt_Token_num=args.Prompt_Token_num, basic_state_dict=basic_state_dict, dec_idx=args.seg_decoder) # setting puzzle_patch_size to not use SAE puzzle_patch_size_scheduler = None fix_position_ratio_scheduler = None # PuzzleTuning elif args.model[0:3] == 'sae': if args.basic_state_dict is not None: try: if args.basic_state_dict == 'timm': print("using timm") basic_model = timm.create_model(basic_encoder + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() else: basic_state_dict = torch.load(args.basic_state_dict) except: print('erro in args.basic_state_dict:', args.basic_state_dict) if args.PromptTuning is not None: print( 'In PromptTuning, the basic_state_dict is required, without specification now, timm loaded.\n') # timm model name basic_encoder basic_model = timm.create_model(basic_encoder + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() else: basic_state_dict = None print('SAE Restart with a empty backbone') else: print('Puzzle tuning with Transfer-learning:', args.basic_state_dict) else: if args.PromptTuning is not None: print('In PromptTuning, the basic_state_dict is required, without specification now, timm loaded.\n') # timm model name basic_encoder basic_model = timm.create_model(basic_encoder + '_' + str(args.input_size), pretrained=True) basic_state_dict = basic_model.state_dict() else: basic_state_dict = None print('Puzzle tuning with a empty backbone') model = SAE.__dict__[args.model](img_size=args.input_size, group_shuffle_size=args.group_shuffle_size, norm_pix_loss=args.norm_pix_loss, prompt_mode=args.PromptTuning, Prompt_Token_num=args.Prompt_Token_num, basic_state_dict=basic_state_dict, dec_idx=args.seg_decoder) fix_position_ratio_scheduler = ratio_scheduler(total_epoches=args.epochs, warmup_epochs=args.warmup_epochs, basic_ratio=0.25, # start ratio fix_position_ratio=args.fix_position_ratio, # None strategy=args.strategy) # strategy=None for fixed else reduce ratio gradually # setting puzzle_patch_size to not use MAE puzzle_patch_size_scheduler = patch_scheduler(total_epoches=args.epochs, warmup_epochs=args.warmup_epochs, edge_size=args.input_size, basic_patch=model.patch_embed.patch_size[0], fix_patch_size=args.fix_patch_size, # None strategy=args.strategy) # 'linear' # NOTICE strategy are used for setting up both the ratio-scheduler and patch-scheduler else: print('This Tuning script only support SAE(PuzzleTuning) or MAE') return -1 # the effective batch size for setting up lr if args.DDP_distributed: eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size() else: eff_batch_size = args.batch_size * torch.cuda.device_count() * args.accum_iter print('eff_batch_size:', eff_batch_size) if args.lr is None: # when only base_lr is specified args.lr = args.blr * eff_batch_size / 256 print("base lr: %.2e" % (args.lr * 256 / eff_batch_size)) print("actual lr: %.2e" % args.lr) print("accumulate grad iterations: %d" % args.accum_iter) print("effective batch size: %d" % eff_batch_size) # take the model parameters for optimizer update model_without_ddp = model if args.DDP_distributed: model.cuda() # args.gpu is obtained by misc.py # find_unused_parameters=True for the DDP to correctly synchronize layers in back propagation model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True) else: model = torch.nn.DataParallel(model) model.to(device) print("Model = %s" % str(model_without_ddp)) # following timm: set wd as 0 for bias and norm layers param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay) optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95)) print(optimizer) # loss scaler with gradient clipping loss_scaler = NativeScaler(GPU_count=torch.cuda.device_count(), DDP_distributed=args.DDP_distributed) # if we have --resume，we will load the checkpoint and continue training, if not, we start a new training # the checkpoint should include model, optimizer, loss_scaler information misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler) # Training by epochs print(f"Start training for {args.epochs} epochs") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): # use args.start_epoch to jump to resume checkpoint if enable_DistributedSampler: # DistributedSampler need to .set_epoch(epoch) at each epoch data_loader_train.sampler.set_epoch(epoch) # training iterations train_stats = train_one_epoch(model, data_loader_train, optimizer, device, epoch, loss_scaler, fix_position_ratio_scheduler=fix_position_ratio_scheduler, puzzle_patch_size_scheduler=puzzle_patch_size_scheduler, check_samples=args.check_samples, print_freq=args.print_freq, log_writer=log_writer, args=args) if args.output_dir and (epoch % args.check_point_gap == 0 or epoch + 1 == args.epochs): misc.save_model(args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, model_idx=args.model_idx) log_stats = {**{f'train_{k}': v for k, v in train_stats.items()}, 'epoch': epoch, } # Write log if args.output_dir and misc.is_main_process(): if log_writer is not None: log_writer.flush() with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f: f.write(json.dumps(log_stats) + "\n") # time stamp total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) def get_args_parser(): parser = argparse.ArgumentParser('SAE pre-training', add_help=False) # disable_notify parser.add_argument('--disable_notify', action='store_true', help='do not send email of tracking') # Model Name or index parser.add_argument('--model_idx', default='PuzzleTuning_', type=str, help='Model Name or index') # Original MAE（224->64) MAE A100（224->256 384->128）SAE（224->128 384->64）SAE-VPT（224->256 384->128） parser.add_argument('--batch_size', default=64, type=int, help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus') parser.add_argument('--epochs', default=200, type=int) # epochs原800 parser.add_argument('--accum_iter', default=2, type=int, help='Accumulate gradient iterations ' '(for increasing the effective batch size under memory constraints)') # if we have --resume，we will load the checkpoint and continue training, if not, we start a new training # the checkpoint should include model, optimizer, loss_scaler information parser.add_argument('--resume', default='', help='resume from checkpoint') parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch of checkpoint') # Model parameters sae_vit_base_patch16 mae_vit_base_patch16 parser.add_argument('--model', default='sae_vit_base_patch16', type=str, metavar='MODEL', help='Name of model to train') # ori mae_vit_large_patch16 parser.add_argument('--seg_decoder', default=None, type=str, metavar='segmentation decoder', help='Name of segmentation decoder') parser.add_argument('--input_size', default=224, type=int,help='images input size') parser.add_argument('--model_patch_size', default=16, type=int, help='model_patch_size, default 16 for ViT-base') parser.add_argument('--num_classes', default=3, type=int, # decoder seg class set to channel help='the number of classes for segmentation') # MAE mask_ratio parser.add_argument('--mask_ratio', default=0.75, type=float, help='Masking ratio (percentage of removed patches)') # Tuning setting # PromptTuning parser.add_argument('--PromptTuning', default=None, type=str, help='use Prompt Tuning strategy (Deep/Shallow) instead of Finetuning (None, by default)') # Prompt_Token_num parser.add_argument('--Prompt_Token_num', default=20, type=int, help='Prompt_Token_num for VPT backbone') # Course learning setting parser.add_argument('--strategy', default=None, type=str, help='use linear or other puzzle size scheduler') parser.add_argument('--fix_position_ratio', default=None, type=float, help='ablation fix_position_ratio (percentage of position token patches)') parser.add_argument('--fix_patch_size', default=None, type=int, help='ablation using fix_patch_size') parser.add_argument('--group_shuffle_size', default=-1, type=int, help='group_shuffle_size of group shuffling,' 'default -1 for the whole batch as a group') # loss settings parser.add_argument('--norm_pix_loss', action='store_true', help='Use (per-patch) normalized pixels as targets for computing loss') parser.set_defaults(norm_pix_loss=False) # basic_state_dict parser.add_argument('--basic_state_dict', default=None, type=str, help='load basic backbone state_dict for Transfer-learning-based tuning, default None') # Optimizer settings parser.add_argument('--weight_decay', type=float, default=0.05, help='weight decay (default: 0.05)') parser.add_argument('--lr', type=float, default=None, metavar='LR', help='learning rate (absolute lr), default=None') parser.add_argument('--blr', type=float, default=1.5e-4, metavar='LR', help='base learning rate: absolute_lr = base_lr * effective batch size / 256') parser.add_argument('--min_lr', type=float, default=0., metavar='LR', help='lower lr bound for cyclic schedulers that hit 0') parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N', help='epochs to warmup LR') # PATH settings # Dataset parameters /datasets01/imagenet_full_size/061417/ /data/imagenet_1k /root/autodl-tmp/imagenet parser.add_argument('--data_path', default='/root/autodl-tmp/datasets/All', type=str, help='dataset path') parser.add_argument('--output_dir', default='/root/autodl-tmp/runs', help='path where to save, empty for no saving') parser.add_argument('--log_dir', default='/root/tf-logs', help='path where to tensorboard log') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=42, type=int) # dataloader setting parser.add_argument('--num_workers', default=20, type=int) # 4A100（16，384，b128, shm40）6A100(36，384，b128, shm100) 8A100(35，384，b128, shm100) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem') parser.set_defaults(pin_mem=True) # print_freq and checkpoint parser.add_argument('--print_freq', default=20, type=int) parser.add_argument('--check_point_gap', default=50, type=int) parser.add_argument('--check_samples', default=1, type=int, help='check how many images in a checking batch') # DDP_distributed training parameters for DDP parser.add_argument('--world_size', default=1, type=int, help='number of DDP_distributed processes') parser.add_argument('--local_rank', default=-1, type=int) parser.add_argument('--dist_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up DDP_distributed training') parser.add_argument('--DDP_distributed', action='store_true', help='Use DDP in training. ' 'without calling, DP with be applied') return parser if __name__ == '__main__': args = get_args_parser() args = args.parse_args() if args.output_dir: Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args)