""" Training Script ver： Oct 23rd 17:30 dataset structure: ImageNet image folder dataset is used. """ from __future__ import print_function, division import argparse import copy import json import time import os import numpy as np import torch import torch.nn as nn import torch.optim as optim import torchvision from tensorboardX import SummaryWriter from torch.optim import lr_scheduler from torchsummary import summary from utils.data_augmentation import data_augmentation from utils.SoftCrossEntropyLoss import SoftlabelCrossEntropy from utils.online_augmentations import get_online_augmentation from utils.visual_usage import visualize_check, check_SAA from utils.tools import setup_seed, del_file, FixStateDict from utils.schedulers import patch_scheduler, ratio_scheduler from Backbone.getmodel import get_model from Backbone.GetPromptModel import build_promptmodel # Training Strategy def better_performance(temp_acc, temp_vac, best_acc, best_vac): # determin which epoch have the best model if temp_vac >= best_vac and temp_acc >= best_acc: return True elif temp_vac > best_vac: return True else: return False def train_model(model, dataloaders, criterion, optimizer, class_names, dataset_sizes, Augmentation=None, fix_position_ratio_scheduler=None, puzzle_patch_size_scheduler=None, edge_size=384, model_idx=None, num_epochs=25, intake_epochs=0, check_minibatch=100, scheduler=None, device=None, draw_path='../imagingresults', enable_attention_check=False, enable_visualize_check=False, enable_sam=False, writer=None): """ Training iteration :param model: model object :param dataloaders: 2 dataloader(train and val) dict :param criterion: loss func obj :param optimizer: optimizer obj :param class_names: The name of classes for priting :param dataset_sizes: size of datasets :param Augmentation: Online augmentation methods :param fix_position_ratio_scheduler: Online augmentation fix_position_ratio_scheduler :param puzzle_patch_size_scheduler: Online augmentation puzzle_patch_size_scheduler :param edge_size: image size for the input image :param model_idx: model idx for the getting pre-setted model :param num_epochs: total training epochs :param intake_epochs: number of skip over epochs when choosing the best model :param check_minibatch: number of skip over minibatch in calculating the criteria's results etc. :param scheduler: scheduler is an LR scheduler object from torch.optim.lr_scheduler. :param device: cpu/gpu object :param draw_path: path folder for output pic :param enable_attention_check: use attention_check to show the pics of models' attention areas :param enable_visualize_check: use visualize_check to show the pics :param enable_sam: use SAM training strategy :param writer: attach the records to the tensorboard backend """ if device is None: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") since = time.time() # for saving the best model state dict best_model_wts = copy.deepcopy(model.state_dict()) # deepcopy # initial an empty dict json_log = {} # initial best performance best_acc = 0.0 best_vac = 0.0 temp_acc = 0.0 temp_vac = 0.0 best_epoch_idx = 1 epoch_loss = 0.0 # initial value for loss-drive for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch + 1, num_epochs)) print('-' * 10) # record json log, initially empty json_log[str(epoch + 1)] = {} # Each epoch has a training and validation phase for phase in ['train', 'val']: # alternatively train/val index = 0 check_index = -1 # set a visulize check at the end of each epoch's train and val model_time = time.time() # initiate the empty log dict log_dict = {} for cls_idx in range(len(class_names)): # only float type is allowed in json, set to float inside log_dict[class_names[cls_idx]] = {'tp': 0.0, 'tn': 0.0, 'fp': 0.0, 'fn': 0.0} if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode # criterias, initially empty running_loss = 0.0 log_running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: # use different dataloder in different phase inputs = inputs.to(device) # print('inputs[0]',type(inputs[0])) # NOTICE in CLS task the labels' type is long tensor([B])，not one-hot ([B,CLS]) labels = labels.to(device) # Online Augmentations on device if Augmentation is not None: if phase == 'train': # cellmix if fix_position_ratio_scheduler is not None and puzzle_patch_size_scheduler is not None: # loss-drive fix_position_ratio = fix_position_ratio_scheduler(epoch, epoch_loss) puzzle_patch_size = puzzle_patch_size_scheduler(epoch, epoch_loss) inputs, labels, GT_long_labels = Augmentation(inputs, labels, fix_position_ratio, puzzle_patch_size) # Counterpart augmentations else: inputs, labels, GT_long_labels = Augmentation(inputs, labels) else: # Val inputs, labels, GT_long_labels = Augmentation(inputs, labels, act=False) else: GT_long_labels = labels # store ori_label on CPU # zero the parameter gradients if not enable_sam: optimizer.zero_grad() # forward # track grad if only in train! with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) # pred outputs of confidence: [B,CLS] _, preds = torch.max(outputs, 1) # idx outputs: [B] each is a idx loss = criterion(outputs, labels) # cross entrphy of one-hot outputs: [B,CLS] and idx label [B] # backward + optimize only if in training phase if phase == 'train': if enable_sam: loss.backward() # first forward-backward pass optimizer.first_step(zero_grad=True) # second forward-backward pass loss2 = criterion(model(inputs), labels) # SAM need another model(inputs) loss2.backward() # make sure to do a full forward pass when using SAM optimizer.second_step(zero_grad=True) else: loss.backward() optimizer.step() # log criterias: update log_running_loss += loss.item() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds.cpu() == GT_long_labels.cpu().data) # Compute precision and recall for each class. for cls_idx in range(len(class_names)): tp = np.dot((GT_long_labels.cpu().data == cls_idx).numpy().astype(int), (preds == cls_idx).cpu().numpy().astype(int)) tn = np.dot((GT_long_labels.cpu().data != cls_idx).numpy().astype(int), (preds != cls_idx).cpu().numpy().astype(int)) fp = np.sum((preds == cls_idx).cpu().numpy()) - tp fn = np.sum((GT_long_labels.cpu().data == cls_idx).numpy()) - tp # log_dict[cls_idx] = {'tp': 0.0, 'tn': 0.0, 'fp': 0.0, 'fn': 0.0} set to float inside log_dict[class_names[cls_idx]]['tp'] += tp log_dict[class_names[cls_idx]]['tn'] += tn log_dict[class_names[cls_idx]]['fp'] += fp log_dict[class_names[cls_idx]]['fn'] += fn # attach the records to the tensorboard backend if writer is not None: # ...log the running loss writer.add_scalar(phase + ' minibatch loss', float(loss.item()), epoch * len(dataloaders[phase]) + index) writer.add_scalar(phase + ' minibatch ACC', float(torch.sum(preds.cpu() == GT_long_labels.cpu().data) / inputs.size(0)), epoch * len(dataloaders[phase]) + 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 epoch_idx = epoch + 1 print('Epoch:', epoch_idx, ' ', phase, 'index of ' + str(check_minibatch) + ' minibatch:', check_index, ' time used:', model_time) print('minibatch AVG loss:', float(log_running_loss) / check_minibatch) if enable_visualize_check: visualize_check(inputs, GT_long_labels, model, class_names, num_images=-1, pic_name='Visual_' + phase + '_E_' + str(epoch_idx) + '_I_' + str(index + 1), draw_path=draw_path, writer=writer) if enable_attention_check: try: check_SAA(inputs, GT_long_labels, model, model_idx, edge_size, class_names, num_images=1, pic_name='GradCAM_' + phase + '_E_' + str(epoch_idx) + '_I_' + str(index + 1), draw_path=draw_path, writer=writer) except: print('model:', model_idx, ' with edge_size', edge_size, 'is not supported yet') else: pass model_time = time.time() log_running_loss = 0.0 index += 1 if phase == 'train': if scheduler is not None: # lr scheduler: update scheduler.step() # at the last of train/val in each epoch, if no check has been triggered if check_index == -1: epoch_idx = epoch + 1 if enable_visualize_check: visualize_check(inputs, GT_long_labels, model, class_names, num_images=-1, pic_name='Visual_' + phase + '_E_' + str(epoch_idx), draw_path=draw_path, writer=writer) if enable_attention_check: try: check_SAA(inputs, GT_long_labels, model, model_idx, edge_size, class_names, num_images=1, pic_name='GradCAM_' + phase + '_E_' + str(epoch_idx), draw_path=draw_path, writer=writer) except: print('model:', model_idx, ' with edge_size', edge_size, 'is not supported yet') else: pass # log criterias: print epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] * 100 print('\nEpoch: {} {} \nLoss: {:.4f} Acc: {:.4f}'.format(epoch + 1, phase, epoch_loss, epoch_acc)) if phase == 'train' and fix_position_ratio_scheduler is not None \ and puzzle_patch_size_scheduler is not None: print('\nEpoch: {}, Fix_position_ratio: {}, Puzzle_patch_size: ' '{}'.format(epoch + 1, fix_position_ratio, puzzle_patch_size)) # attach the records to the tensorboard backend if writer is not None: # ...log the running loss writer.add_scalar(phase + ' loss', float(epoch_loss), epoch + 1) writer.add_scalar(phase + ' ACC', float(epoch_acc), epoch + 1) # calculating the confusion matrix for cls_idx in range(len(class_names)): tp = log_dict[class_names[cls_idx]]['tp'] tn = log_dict[class_names[cls_idx]]['tn'] fp = log_dict[class_names[cls_idx]]['fp'] fn = log_dict[class_names[cls_idx]]['fn'] tp_plus_fp = tp + fp tp_plus_fn = tp + fn fp_plus_tn = fp + tn fn_plus_tn = fn + tn # precision if tp_plus_fp == 0: precision = 0 else: precision = float(tp) / tp_plus_fp * 100 # recall if tp_plus_fn == 0: recall = 0 else: recall = float(tp) / tp_plus_fn * 100 # TPR (sensitivity) TPR = recall # TNR (specificity) # FPR if fp_plus_tn == 0: TNR = 0 FPR = 0 else: TNR = tn / fp_plus_tn * 100 FPR = fp / fp_plus_tn * 100 # NPV if fn_plus_tn == 0: NPV = 0 else: NPV = tn / fn_plus_tn * 100 print('{} precision: {:.4f} recall: {:.4f}'.format(class_names[cls_idx], precision, recall)) print('{} sensitivity: {:.4f} specificity: {:.4f}'.format(class_names[cls_idx], TPR, TNR)) print('{} FPR: {:.4f} NPV: {:.4f}'.format(class_names[cls_idx], FPR, NPV)) print('{} TP: {}'.format(class_names[cls_idx], tp)) print('{} TN: {}'.format(class_names[cls_idx], tn)) print('{} FP: {}'.format(class_names[cls_idx], fp)) print('{} FN: {}'.format(class_names[cls_idx], fn)) # attach the records to the tensorboard backend if writer is not None: # ...log the running loss writer.add_scalar(phase + ' ' + class_names[cls_idx] + ' precision', precision, epoch + 1) writer.add_scalar(phase + ' ' + class_names[cls_idx] + ' recall', recall, epoch + 1) # json log: update json_log[str(epoch + 1)][phase] = log_dict if phase == 'val': temp_vac = epoch_acc else: temp_acc = epoch_acc # not useful actually # deep copy the model if phase == 'val' and better_performance(temp_acc, temp_vac, best_acc, best_vac) and epoch >= intake_epochs: # what is better? we now use the wildly used method only best_epoch_idx = epoch + 1 best_acc = temp_acc best_vac = temp_vac best_model_wts = copy.deepcopy(model.state_dict()) best_log_dic = log_dict print('\n') print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) print('Best epoch idx: ', best_epoch_idx) print('Best epoch train Acc: {:4f}'.format(best_acc)) print('Best epoch val Acc: {:4f}'.format(best_vac)) for cls_idx in range(len(class_names)): tp = best_log_dic[class_names[cls_idx]]['tp'] tn = best_log_dic[class_names[cls_idx]]['tn'] fp = best_log_dic[class_names[cls_idx]]['fp'] fn = best_log_dic[class_names[cls_idx]]['fn'] tp_plus_fp = tp + fp tp_plus_fn = tp + fn fp_plus_tn = fp + tn fn_plus_tn = fn + tn # precision if tp_plus_fp == 0: precision = 0 else: precision = float(tp) / tp_plus_fp * 100 # recall if tp_plus_fn == 0: recall = 0 else: recall = float(tp) / tp_plus_fn * 100 # TPR (sensitivity) TPR = recall # TNR (specificity) # FPR if fp_plus_tn == 0: TNR = 0 FPR = 0 else: TNR = tn / fp_plus_tn * 100 FPR = fp / fp_plus_tn * 100 # NPV if fn_plus_tn == 0: NPV = 0 else: NPV = tn / fn_plus_tn * 100 print('{} precision: {:.4f} recall: {:.4f}'.format(class_names[cls_idx], precision, recall)) print('{} sensitivity: {:.4f} specificity: {:.4f}'.format(class_names[cls_idx], TPR, TNR)) print('{} FPR: {:.4f} NPV: {:.4f}'.format(class_names[cls_idx], FPR, NPV)) # attach the records to the tensorboard backend if writer is not None: writer.close() # load best model weights as final model training result model.load_state_dict(best_model_wts) # save json_log indent=2 for better view json.dump(json_log, open(os.path.join(draw_path, model_idx + '_log.json'), 'w'), ensure_ascii=False, indent=2) return model def main(args): if args.paint: # use Agg kernal, not painting in the front-desk import matplotlib matplotlib.use('Agg') enable_tensorboard = args.enable_tensorboard # True enable_attention_check = args.enable_attention_check # False 'CAM' 'SAA' enable_visualize_check = args.enable_visualize_check # False enable_sam = args.enable_sam # False data_augmentation_mode = args.data_augmentation_mode # 0 linearprobing = args.linearprobing # False Pre_Trained_model_path = args.Pre_Trained_model_path # None Prompt_state_path = args.Prompt_state_path # None # Prompt PromptTuning = args.PromptTuning # None "Deep" / "Shallow" Prompt_Token_num = args.Prompt_Token_num # 20 PromptUnFreeze = args.PromptUnFreeze # False gpu_idx = args.gpu_idx # GPU idx start with0, -1 to use multipel GPU # model info model_idx = args.model_idx # the model we are going to use. by the format of Model_size_other_info # structural parameter drop_rate = args.drop_rate attn_drop_rate = args.attn_drop_rate drop_path_rate = args.drop_path_rate use_cls_token = False if args.cls_token_off else True use_pos_embedding = False if args.pos_embedding_off else True use_att_module = None if args.att_module == 'None' else args.att_module # pretrained_backbone pretrained_backbone = False if args.backbone_PT_off else True # classification required number of your dataset num_classes = args.num_classes # default 0 for auto-fit # image size for the input image edge_size = args.edge_size # 224 384 1000 # batch info batch_size = args.batch_size # 8 num_workers = args.num_workers # main training num_workers 4 num_epochs = args.num_epochs # 50 intake_epochs = args.intake_epochs # 0 check_minibatch = args.check_minibatch if args.check_minibatch is not None else 400 // batch_size lr = args.lr # 0.000007 lrf = args.lrf # 0.0 opt_name = args.opt_name # 'Adam' # PATH info draw_root = args.draw_root model_path = args.model_path dataroot = args.dataroot draw_path = os.path.join(draw_root, 'CLS_' + model_idx) # CLS_ is for the CLS training, MIL will be MIL training save_model_path = os.path.join(model_path, 'CLS_' + model_idx + '.pth') if not os.path.exists(model_path): os.makedirs(model_path) if os.path.exists(draw_path): del_file(draw_path) # fixme clear the output folder, NOTICE this may be DANGEROUS else: os.makedirs(draw_path) # Train Augmentation augmentation_name = args.augmentation_name # None # Data Augmentation data_transforms = data_augmentation(data_augmentation_mode, edge_size=edge_size) datasets = {x: torchvision.datasets.ImageFolder(os.path.join(dataroot, x), data_transforms[x]) for x in ['train', 'val']} # 2 dataset obj is prepared here and combine together dataset_sizes = {x: len(datasets[x]) for x in ['train', 'val']} # size of each dataset dataloaders = {'train': torch.utils.data.DataLoader(datasets['train'], batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True), # colab suggest 2 'val': torch.utils.data.DataLoader(datasets['val'], batch_size=batch_size, shuffle=False, num_workers=num_workers // 4 + 1, drop_last=True) } class_names = [d.name for d in os.scandir(os.path.join(dataroot, 'train')) if d.is_dir()] class_names.sort() if num_classes == 0: print("class_names:", class_names) num_classes = len(class_names) else: if len(class_names) == num_classes: print("class_names:", class_names) else: print('classfication number of the model mismatch the dataset requirement of:', len(class_names)) return -1 print("*********************************{}*************************************".format('setting')) print(args) # start tensorboard backend if enable_tensorboard: writer = SummaryWriter(draw_path) else: writer = None # if u run locally # nohup tensorboard --logdir=/home/MSHT/runs --host=0.0.0.0 --port=7777 & # tensorboard --logdir=/home/ZTY/runs --host=0.0.0.0 --port=7777 if gpu_idx == -1: # use all cards if torch.cuda.device_count() > 1: print("Use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs gpu_use = gpu_idx else: print('we dont have more GPU idx here, try to use gpu_idx=0') try: os.environ['CUDA_VISIBLE_DEVICES'] = '0' # setting k for: only card idx k is sighted for this code gpu_use = 0 except: print("GPU distributing ERRO occur use CPU instead") gpu_use = 'cpu' else: # Decide which device we want to run on try: # setting k for: only card idx k is sighted for this code os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_idx) gpu_use = gpu_idx except: print('we dont have that GPU idx here, try to use gpu_idx=0') try: # setting 0 for: only card idx 0 is sighted for this code os.environ['CUDA_VISIBLE_DEVICES'] = '0' gpu_use = 0 except: print("GPU distributing ERRO occur use CPU instead") gpu_use = 'cpu' # device environment device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # get model if PromptTuning is not None: print('PromptTuning of ', model_idx) print('Prompt VPT type:', PromptTuning) # initialize the model backbone: if Pre_Trained_model_path is None or Pre_Trained_model_path == 'timm': base_state_dict = 'timm' print('backbone base_state_dict of timm') elif Pre_Trained_model_path is not None and os.path.exists(Pre_Trained_model_path): print('backbone base_state_dict at: ', Pre_Trained_model_path) base_state_dict = torch.load(Pre_Trained_model_path) else: print('invalid Pre_Trained_model_path for prompting at: ', Pre_Trained_model_path) raise # put the additional prompt tokens to model: if Prompt_state_path is None: prompt_state_dict = None print('prompting with empty prompt_state: prompt_state of None') elif Prompt_state_path is not None and os.path.exists(Prompt_state_path): print('prompting with prompt_state at: ', Prompt_state_path) prompt_state_dict = torch.load(Prompt_state_path) else: print('invalid prompt_state_dict for prompting, path at:', Prompt_state_path) raise model = build_promptmodel(num_classes, edge_size, model_idx, Prompt_Token_num=Prompt_Token_num, VPT_type=PromptTuning, prompt_state_dict=prompt_state_dict, base_state_dict=base_state_dict) # Use FineTuning with prompt tokens (when PromptUnFreeze==True) if PromptUnFreeze: model.UnFreeze() print('prompt tuning with all parameaters un-freezed') else: # get model: randomly initiate model, except the backbone CNN(when pretrained_backbone is True) model = get_model(num_classes, edge_size, model_idx, drop_rate, attn_drop_rate, drop_path_rate, pretrained_backbone, use_cls_token, use_pos_embedding, use_att_module) # Manually get the model pretrained on the Imagenet1000 if Pre_Trained_model_path is not None: if os.path.exists(Pre_Trained_model_path): state_dict = FixStateDict(torch.load(Pre_Trained_model_path), remove_key_head='head') model.load_state_dict(state_dict, False) print('Specified backbone model weight loaded:', Pre_Trained_model_path) else: print('Specified Pre_Trained_model_path:' + Pre_Trained_model_path, ' is NOT avaliable!!!!\n') raise else: print('building model (no-prompt) with pretrained_backbone status:',pretrained_backbone) if pretrained_backbone is True: print('timm loaded') if linearprobing: # Only tuning the last FC layer for CLS task module_all = 0 for child in model.children(): # find all nn.modules module_all += 1 for param in model.parameters(): # freeze all parameters param.requires_grad = False for module_idx, child in enumerate(model.children()): if module_idx == module_all: # Unfreeze the parameters of the last FC layer for param in child.parameters(): param.requires_grad = True print('GPU:', gpu_use) if gpu_use == -1: model = nn.DataParallel(model) model.to(device) try: summary(model, input_size=(3, edge_size, edge_size)) # should be after .to(device) except: pass print("model :", model_idx) # Augmentation Augmentation = get_online_augmentation(augmentation_name, p=0.5, class_num=num_classes, batch_size=batch_size, edge_size=edge_size, device=device) if augmentation_name != 'CellMix-Split' and augmentation_name != 'CellMix-Group' \ and augmentation_name != 'CellMix-Random': fix_position_ratio_scheduler = None puzzle_patch_size_scheduler = None else: # setting puzzle_patch_size and fix_position_ratio schedulers fix_position_ratio_scheduler = ratio_scheduler(total_epoches=num_epochs, warmup_epochs=0, basic_ratio=0.5, strategy=args.ratio_strategy, # 'linear' fix_position_ratio=args.fix_position_ratio, threshold=args.loss_drive_threshold) puzzle_patch_size_scheduler = patch_scheduler(total_epoches=num_epochs, warmup_epochs=0, edge_size=edge_size, basic_patch=16, strategy=args.patch_strategy, # 'random', 'linear' or 'loop' threshold=args.loss_drive_threshold, fix_patch_size=args.fix_patch_size, # 16,32,48,64,96,128,192 patch_size_jump=args.patch_size_jump) # 'odd' or 'even' # Default cross entrphy of one-hot outputs: [B,CLS] and idx label [B] long tensor # augmentation loss is SoftlabelCrossEntropy criterion = SoftlabelCrossEntropy() \ if Augmentation is not None and augmentation_name != 'Cutout' else nn.CrossEntropyLoss() if opt_name == 'SGD': optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.8, weight_decay=0.005) scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5) # 15 0.1 default SGD StepLR scheduler elif opt_name == 'Adam': optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.01) scheduler = None else: print('no optimizer') raise if enable_sam: from utils.sam import SAM if opt_name == 'SGD': base_optimizer = torch.optim.SGD # define an optimizer for the "sharpness-aware" update optimizer = SAM(model.parameters(), base_optimizer, lr=lr, momentum=0.8) scheduler = None elif opt_name == 'Adam': base_optimizer = torch.optim.Adam # define an optimizer for the "sharpness-aware" update optimizer = SAM(model.parameters(), base_optimizer, lr=lr, weight_decay=0.01) else: print('no optimizer') raise if lrf > 0: # use cosine learning rate schedule import math # cosine Scheduler by https://arxiv.org/pdf/1812.01187.pdf lf = lambda x: ((1 + math.cos(x * math.pi / num_epochs)) / 2) * (1 - lrf) + lrf # cosine scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # train model_ft = train_model(model, dataloaders, criterion, optimizer, class_names, dataset_sizes, fix_position_ratio_scheduler=fix_position_ratio_scheduler, puzzle_patch_size_scheduler=puzzle_patch_size_scheduler, Augmentation=Augmentation, edge_size=edge_size, model_idx=model_idx, num_epochs=num_epochs, intake_epochs=intake_epochs, check_minibatch=check_minibatch, scheduler=scheduler, device=device, draw_path=draw_path, enable_attention_check=enable_attention_check, enable_visualize_check=enable_visualize_check, enable_sam=enable_sam, writer=writer) # save model if its a multi-GPU model, save as a single GPU one too if gpu_use == -1: if PromptTuning is None: torch.save(model_ft.module.state_dict(), save_model_path) else: if PromptUnFreeze: torch.save(model_ft.module.state_dict(), save_model_path) else: prompt_state_dict = model_ft.module.obtain_prompt() # fixme maybe bug at DP module.obtain_prompt, just model.obtain_prompt is enough torch.save(prompt_state_dict, save_model_path) print('model trained by multi-GPUs has its single GPU copy saved at ', save_model_path) else: if PromptTuning is None: torch.save(model_ft.state_dict(), save_model_path) else: if PromptUnFreeze: torch.save(model_ft.state_dict(), save_model_path) else: prompt_state_dict = model_ft.obtain_prompt() torch.save(prompt_state_dict, save_model_path) print('model trained by GPU (idx:' + str(gpu_use) + ') has been saved at ', save_model_path) def get_args_parser(): parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') # Model Name or index parser.add_argument('--model_idx', default='Hybrid2_384_401_testsample', type=str, help='Model Name or index') # drop_rate, attn_drop_rate, drop_path_rate parser.add_argument('--drop_rate', default=0.0, type=float, help='dropout rate , default 0.0') parser.add_argument('--attn_drop_rate', default=0.0, type=float, help='dropout rate Aftter Attention, default 0.0') parser.add_argument('--drop_path_rate', default=0.0, type=float, help='drop path for stochastic depth, default 0.0') # Abalation Studies parser.add_argument('--cls_token_off', action='store_true', help='use cls_token in model structure') parser.add_argument('--pos_embedding_off', action='store_true', help='use pos_embedding in model structure') # 'SimAM', 'CBAM', 'SE' 'None' parser.add_argument('--att_module', default='SimAM', type=str, help='use which att_module in model structure') # backbone_PT_off by default is false, in default setting the backbone weight is required parser.add_argument('--backbone_PT_off', action='store_true', help='use a freash backbone weight in training') # Enviroment parameters parser.add_argument('--gpu_idx', default=-1, type=int, help='use a single GPU with its index, -1 to use multiple GPU') # Path parameters parser.add_argument('--dataroot', default='/data/MIL_Experiment/dataset/ROSE_CLS', help='path to dataset') parser.add_argument('--model_path', default='/home/pancreatic-cancer-project/saved_models', help='path to save model state-dict') parser.add_argument('--draw_root', default='/home/pancreatic-cancer-project/runs', help='path to draw and save tensorboard output') # Help tool parameters parser.add_argument('--paint', action='store_false', help='paint in front desk') # matplotlib.use('Agg') # check tool parameters parser.add_argument('--enable_tensorboard', action='store_true', help='enable tensorboard to save status') parser.add_argument('--enable_attention_check', action='store_true', help='check and save attention map') parser.add_argument('--enable_visualize_check', action='store_true', help='check and save pics') # Tuning setting # PromptTuning parser.add_argument('--PromptTuning', default=None, type=str, help='use Prompt Tuning strategy instead of Finetuning') # Prompt_Token_num parser.add_argument('--Prompt_Token_num', default=20, type=int, help='Prompt_Token_num') # PromptUnFreeze parser.add_argument('--PromptUnFreeze', action='store_true', help='prompt tuning with all parameaters un-freezed') # linearprobing parser.add_argument('--linearprobing', action='store_true', help='use linearprobing tuning') # Finetuning a Pretrained model at PATH # '/home/MIL_Experiment/saved_models/Hybrid2_384_PreTrain_000.pth' parser.add_argument('--Pre_Trained_model_path', default=None, type=str, help='Finetuning a trained model in this dataset') # Prompt_state_path parser.add_argument('--Prompt_state_path', default=None, type=str, help='Prompt_state_path for prompt tokens') # Training status parameters # SAM parser.add_argument('--enable_sam', action='store_true', help='use SAM strategy in training') # Online augmentation_name parser.add_argument('--augmentation_name', default=None, type=str, help='Online augmentation name') # CellMix ablation: loss_drive strategy parser.add_argument('--ratio_strategy', default=None, type=str, help='CellMix ratio scheduler strategy') parser.add_argument('--patch_strategy', default=None, type=str, help='CellMix patch scheduler strategy') parser.add_argument('--loss_drive_threshold', default=4.0, type=float, help='CellMix loss_drive_threshold') # CellMix ablation: fix_patch_size patch_size_jump parser.add_argument('--fix_position_ratio', default=0.5, type=float, help='CellMix ratio scheduler strategy') parser.add_argument('--fix_patch_size', default=None, type=int, help='CellMix ablation using fix_patch_size') parser.add_argument('--patch_size_jump', default=None, type=str, help='CellMix patch_size_jump strategy') # Dataset based parameters parser.add_argument('--num_classes', default=0, type=int, help='classification number, default 0 for auto-fit') parser.add_argument('--edge_size', default=384, type=int, help='edge size of input image') # 224 256 384 1000 # Dataset specific augmentations in dataloader parser.add_argument('--data_augmentation_mode', default=0, type=int, help='data_augmentation_mode') # Training seting parameters parser.add_argument('--batch_size', default=8, type=int, help='Training batch_size default 8') parser.add_argument('--num_epochs', default=50, type=int, help='training epochs') parser.add_argument('--intake_epochs', default=0, type=int, help='only save model at epochs after intake_epochs') parser.add_argument('--lr', default=0.00001, type=float, help='learing rate') parser.add_argument('--lrf', type=float, default=0.0, help='learing rate decay rate, default 0(not enabled), suggest 0.1 and lr=0.00005') parser.add_argument('--opt_name', default='Adam', type=str, help='optimizer name Adam or SGD') # check_minibatch for painting pics parser.add_argument('--check_minibatch', default=None, type=int, help='check batch_size') parser.add_argument('--num_workers', default=2, type=int, help='use CPU num_workers , default 2 for colab') return parser if __name__ == '__main__': # setting up the random seed setup_seed(42) parser = get_args_parser() args = parser.parse_args() main(args)