""" Testing script of PuzzleTuning Visualization 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 independent testing for step 2 update: Use "--seg_decoder" parameter to introduce segmentation networks swin_unet for Swin-Unet """ import argparse import datetime 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 from SSL_structures import models_mae, SAE from utils.visual_usage import patchify, unpatchify, Draw_tri_fig from torchvision.transforms import ToPILImage def Puzzle_test(model, data_loader_test, test_dataset_size, mask_ratio, fix_position_ratio, fix_patch_size, check_minibatch=100, enable_visualize_check=True, combined_pred_illustration=False, check_samples=1, device=None, output_dir=None, writer=None, args=None): # start testing print(f"Start testing for {args.model_idx} \n with checkpoint: {args.checkpoint_path}") start_time = time.time() index = 0 model_time = time.time() # criterias, initially empty running_loss = 0.0 log_running_loss = 0.0 model.eval() # Iterate over data. for inputs, labels in data_loader_test: # use different dataloder in different phase inputs = inputs.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) # for tracking fixme if args.model[0:3] == 'sae': loss, pred, imgs_puzzled_patches = model(inputs, fix_position_ratio=fix_position_ratio, puzzle_patch_size=fix_patch_size, combined_pred_illustration=combined_pred_illustration) # SAE else: # args.model[0:3] == 'mae' loss, pred, mask_patch_indicators = model(inputs, mask_ratio=mask_ratio) # MAE loss_value = float(loss.cpu().detach().numpy()) if args.gpu == 1 else sum(loss.cpu().detach().numpy()) # log criterias: update log_running_loss += loss_value running_loss += loss_value * inputs.size(0) # attach the records to the tensorboard backend if writer is not None: # ...log the running loss writer.add_scalar('Test minibatch loss', float(loss_value), index) # at the checking time now if index % check_minibatch == check_minibatch - 1: model_time = time.time() - model_time check_index = index // check_minibatch + 1 print('Test index ' + str(check_index) + ' of ' + str(check_minibatch) + ' minibatch with batch_size of ' + str(inputs.size(0)) + ' time used:', model_time) print('minibatch AVG loss:', float(log_running_loss) / check_minibatch) model_time = time.time() log_running_loss = 0.0 # paint pic if enable_visualize_check: if args.model[0:3] == 'sae': imgs_puzzled_batch = unpatchify(imgs_puzzled_patches, patch_size=16) # Reconstructed img recons_img_batch = unpatchify(pred, patch_size=16) else: # MAE sample_img_patches = patchify(inputs, patch_size=16) # on GPU masked_img_patches = sample_img_patches * mask_patch_indicators.unsqueeze(-1).expand(-1, -1, sample_img_patches.shape[-1]) masked_img_batch = unpatchify(masked_img_patches, patch_size=16) if combined_pred_illustration: anti_mask_patch_indicators = 1 - mask_patch_indicators pred_img_patches = pred * anti_mask_patch_indicators.unsqueeze(-1).\ expand(-1, -1, sample_img_patches.shape[-1]) # Reconstructed img recons_img_batch = unpatchify(masked_img_patches + pred_img_patches, patch_size=16) else: # Reconstructed img recons_img_batch = unpatchify(pred, patch_size=16) for sampleIDX in range(check_samples): # Ori img sample_img = inputs.cpu()[sampleIDX] sample_img = ToPILImage()(sample_img) sample_img.save(os.path.join(output_dir, 'Test_sample_idx_' + str(check_index) + '_sampleIDX_' + str(sampleIDX) + '.jpg')) recons_img = recons_img_batch.cpu()[sampleIDX] recons_img = ToPILImage()(recons_img) recons_img.save(os.path.join(output_dir, 'Test_recons_idx_' + str(check_index) + '_sampleIDX_' + str(sampleIDX) + '.jpg')) # mask_img or puzzled_img if args.model[0:3] == 'sae': puzzled_img = imgs_puzzled_batch.cpu()[sampleIDX] puzzled_img = ToPILImage()(puzzled_img) puzzled_img.save(os.path.join(output_dir, 'Test_puzzled_idx_' + str(check_index) + '.jpg')) picpath = os.path.join(output_dir, 'Test_minibatchIDX_' + str(check_index) + '_sampleIDX_' + str(sampleIDX) + '.jpg') Draw_tri_fig(sample_img, puzzled_img, recons_img, picpath) else: # MAE masked_img = masked_img_batch.cpu()[sampleIDX] masked_img = ToPILImage()(masked_img) masked_img.save(os.path.join(output_dir, 'Test_masked_idx_' + str(check_index) + '.jpg')) picpath = os.path.join(output_dir, 'Test_minibatchIDX_' + str(check_index) + '_sampleIDX_' + str(sampleIDX) + '.jpg') Draw_tri_fig(sample_img, masked_img, recons_img, picpath) index += 1 # time stamp total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) # log criterias: print epoch_loss = running_loss / test_dataset_size print('\nTest_dataset_size: {} \nAvg Loss: {:.4f}'.format(test_dataset_size, epoch_loss)) print('Testing time {}'.format(total_time_str)) def main(args): # 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 # Specify the Test settings if args.fix_position_ratio is not None and args.fix_patch_size is not None and args.mask_ratio is None: args.model_idx = 'Testing_' + args.model_idx + '_b_' + str(args.batch_size) \ + '_hint_ratio_' + str(args.fix_position_ratio) + '_patch_size_' + str(args.fix_patch_size) elif args.mask_ratio is not None and args.fix_position_ratio is None and args.fix_patch_size is None: args.model_idx = 'Testing_' + args.model_idx + '_b_' + str(args.batch_size) \ + '_mask_ratio_' + str(args.mask_ratio) else: print('not a correct test setting, should correctly specify fix_position_ratio/fix_patch_size/mask_ratio') print('\n\n' + args.model_idx + '\n\n') # setting k for: only card idx k is sighted for this code if args.gpu_idx != -1: # fixme: notice for test, we are going to use single gpu only print("Use", torch.cuda.device_count(), "GPUs of idx:", args.gpu_idx) os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_idx) else: print("Use", torch.cuda.device_count(), "GPUs") args.gpu = torch.cuda.device_count() print('job AImageFolderDir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) print("{}".format(args).replace(', ', ',\n')) device = torch.device(args.device) # cuda # fix the seed for reproducibility torch.manual_seed(args.seed) np.random.seed(args.seed) cudnn.benchmark = True # simple augmentation transform_test = 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.ToTensor(), # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) test_dataroot = os.path.join(args.data_path) # , 'test' dataset_test = datasets.ImageFolder(test_dataroot, transform=transform_test) test_dataset_size = len(dataset_test) class_names = [d.name for d in os.scandir(test_dataroot) if d.is_dir()] class_names.sort() print('dataset_test', dataset_test) # Test data # skip minibatch, none to draw 80 figs check_minibatch = args.check_minibatch if args.check_minibatch is not None else \ test_dataset_size // (80 * args.batch_size) check_minibatch = max(1, check_minibatch) # outputs if 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 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) print('Testing output files will be at', args.output_dir) data_loader_test = torch.utils.data.DataLoader(dataset_test, shuffle=args.shuffle_dataloader, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, # 建议False drop_last=True) # define the model if args.model[0:3] == 'mae': 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, dec_idx=args.seg_decoder) elif args.model[0:3] == 'sae': 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, dec_idx=args.seg_decoder) else: print('This MIM test script only support SAE or MAE') return -1 # take model out of checkpoint and load_model state_dict = torch.load(args.checkpoint_path)['model'] model.load_state_dict(state_dict, False) model.to(device) # loss backward and optimizer operations and no longer needed in testing # loss_scaler = NativeScaler() Puzzle_test(model, data_loader_test, test_dataset_size, args.mask_ratio, args.fix_position_ratio, args.fix_patch_size, check_minibatch, args.enable_visualize_check, args.combined_pred_illustration, args.check_samples, device=device, output_dir=args.output_dir, writer=log_writer, args=args) # os.system("shutdown") # AUTO-DL server shutdown currently moved to .sh script for nohup task queue. def get_args_parser(): parser = argparse.ArgumentParser('MIM visualization for PuzzleTuning', add_help=False) # Model Name or index parser.add_argument('--model_idx', default='PuzzleTuning_', type=str, help='Model Name or index') # testing batch size parser.add_argument('--batch_size', default=16, type=int, help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus') # Model parameters sae_vit_base_patch16 or 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, # 原224 help='images input size') 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=None, type=float, help='Masking ratio (percentage of removed patches).') # Hint tokens parser.add_argument('--fix_position_ratio', default=None, type=float, help='basic fix_position_ratio (percentage of position token patches).') parser.add_argument('--fix_patch_size', default=None, type=int, # 原224 help='images input 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') # shuffle_dataloader parser.add_argument('--shuffle_dataloader', action='store_true', help='shuffle Test dataset') # Tuning setting # PromptTuning parser.add_argument('--PromptTuning', default=None, type=str, help='Deep/Shallow to use Prompt Tuning model instead of Finetuning model, by default None') # Prompt_Token_num parser.add_argument('--Prompt_Token_num', default=20, type=int, help='Prompt_Token_num') # 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) # PATH settings # Dataset parameters /root/autodl-tmp/MARS_ALL /root/autodl-tmp/imagenet /root/autodl-tmp/datasets/All parser.add_argument('--data_path', default='/root/autodl-tmp/datasets/PuzzleTuning_demoset', type=str, help='dataset path') parser.add_argument('--output_dir', default='/root/autodl-tmp/runs', help='path where to save test log, empty for no saving') parser.add_argument('--log_dir', default='/root/tf-logs', help='path where to test tensorboard log') # Enviroment parameters parser.add_argument('--gpu_idx', default=0, type=int, help='use a single GPU with its index, -1 to use multiple GPU') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=42, type=int) # ori 0 不过应该无所谓？ # checkpoint_state_dict_path parser.add_argument('--checkpoint_path', default='/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_Prompt_Deep_tokennum_20_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth', type=str, help='load state_dict for testing') # check settings parser.add_argument('--combined_pred_illustration', action='store_true', help='check combined_pred_illustration pics') parser.add_argument('--enable_visualize_check', action='store_true', help='check and save pics') parser.add_argument('--check_minibatch', default=None, type=int, help='check batch_size') parser.add_argument('--check_samples', default=1, type=int, help='check how many images in a checking batch') # dataloader setting parser.add_argument('--num_workers', default=10, type=int) 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) 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)