Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import numpy as np | |
| import datetime | |
| import logging | |
| import json | |
| import torch | |
| import shutil | |
| def calculateNorm2(model): | |
| para_norm = 0. | |
| for p in model.parameters(): | |
| para_norm += p.data.norm(2) | |
| print('2-norm of the neural network: {:.4f}'.format(para_norm**.5)) | |
| def showLR(optimizer): | |
| return optimizer.param_groups[0]['lr'] | |
| class AverageMeter(object): | |
| """Computes and stores the average and current value""" | |
| def __init__(self): | |
| self.reset() | |
| def reset(self): | |
| self.val = 0 | |
| self.avg = 0 | |
| self.sum = 0 | |
| self.count = 0 | |
| def update(self, val, n=1): | |
| self.val = val | |
| self.sum += val * n | |
| self.count += n | |
| self.avg = self.sum / self.count | |
| # -- IO utils | |
| def read_txt_lines(filepath): | |
| assert os.path.isfile( filepath ), "Error when trying to read txt file, path does not exist: {}".format(filepath) | |
| with open( filepath ) as myfile: | |
| content = myfile.read().splitlines() | |
| return content | |
| def save_as_json(d, filepath): | |
| with open(filepath, 'w') as outfile: | |
| json.dump(d, outfile, indent=4, sort_keys=True) | |
| def load_json( json_fp ): | |
| assert os.path.isfile( json_fp ), "Error loading JSON. File provided does not exist, cannot read: {}".format( json_fp ) | |
| with open( json_fp, 'r' ) as f: | |
| json_content = json.load(f) | |
| return json_content | |
| def save2npz(filename, data=None): | |
| assert data is not None, "data is {}".format(data) | |
| if not os.path.exists(os.path.dirname(filename)): | |
| os.makedirs(os.path.dirname(filename)) | |
| np.savez_compressed(filename, data=data) | |
| # -- checkpoints | |
| class CheckpointSaver: | |
| def __init__(self, save_dir, checkpoint_fn='ckpt.pth.tar', best_fn='ckpt.best.pth.tar', best_step_fn='ckpt.best.step{}.pth.tar', save_best_step=False, lr_steps=[]): | |
| """ | |
| Only mandatory: save_dir | |
| Can configure naming of checkpoint files through checkpoint_fn, best_fn and best_stage_fn | |
| If you want to keep best-performing checkpoint per step | |
| """ | |
| self.save_dir = save_dir | |
| # checkpoint names | |
| self.checkpoint_fn = checkpoint_fn | |
| self.best_fn = best_fn | |
| self.best_step_fn = best_step_fn | |
| # save best per step? | |
| self.save_best_step = save_best_step | |
| self.lr_steps = [] | |
| # init var to keep track of best performing checkpoint | |
| self.current_best = 0 | |
| # save best at each step? | |
| if self.save_best_step: | |
| assert lr_steps != [], "Since save_best_step=True, need proper value for lr_steps. Current: {}".format(lr_steps) | |
| self.best_for_stage = [0]*(len(lr_steps)+1) | |
| def save(self, save_dict, current_perf, epoch=-1): | |
| """ | |
| Save checkpoint and keeps copy if current perf is best overall or [optional] best for current LR step | |
| """ | |
| # save last checkpoint | |
| checkpoint_fp = os.path.join(self.save_dir, self.checkpoint_fn) | |
| # keep track of best model | |
| self.is_best = current_perf > self.current_best | |
| if self.is_best: | |
| self.current_best = current_perf | |
| best_fp = os.path.join(self.save_dir, self.best_fn) | |
| save_dict['best_prec'] = self.current_best | |
| # keep track of best-performing model per step [optional] | |
| if self.save_best_step: | |
| assert epoch >= 0, "Since save_best_step=True, need proper value for 'epoch'. Current: {}".format(epoch) | |
| s_idx = sum( epoch >= l for l in lr_steps ) | |
| self.is_best_for_stage = current_perf > self.best_for_stage[s_idx] | |
| if self.is_best_for_stage: | |
| self.best_for_stage[s_idx] = current_perf | |
| best_stage_fp = os.path.join(self.save_dir, self.best_stage_fn.format(s_idx)) | |
| save_dict['best_prec_per_stage'] = self.best_for_stage | |
| # save | |
| torch.save(save_dict, checkpoint_fp) | |
| print("Checkpoint saved at {}".format(checkpoint_fp)) | |
| if self.is_best: | |
| shutil.copyfile(checkpoint_fp, best_fp) | |
| if self.save_best_step and self.is_best_for_stage: | |
| shutil.copyfile(checkpoint_fp, best_stage_fp) | |
| def set_best_from_ckpt(self, ckpt_dict): | |
| self.current_best = ckpt_dict['best_prec'] | |
| self.best_for_stage = ckpt_dict.get('best_prec_per_stage',None) | |
| def load_model(load_path, model, optimizer = None, allow_size_mismatch = False): | |
| """ | |
| Load model from file | |
| If optimizer is passed, then the loaded dictionary is expected to contain also the states of the optimizer. | |
| If optimizer not passed, only the model weights will be loaded | |
| """ | |
| # -- load dictionary | |
| assert os.path.isfile( load_path ), "Error when loading the model, provided path not found: {}".format( load_path ) | |
| checkpoint = torch.load(load_path) | |
| loaded_state_dict = checkpoint['model_state_dict'] | |
| if allow_size_mismatch: | |
| loaded_sizes = { k: v.shape for k,v in loaded_state_dict.items() } | |
| model_state_dict = model.state_dict() | |
| model_sizes = { k: v.shape for k,v in model_state_dict.items() } | |
| mismatched_params = [] | |
| for k in loaded_sizes: | |
| if loaded_sizes[k] != model_sizes[k]: | |
| mismatched_params.append(k) | |
| for k in mismatched_params: | |
| del loaded_state_dict[k] | |
| # -- copy loaded state into current model and, optionally, optimizer | |
| model.load_state_dict(loaded_state_dict, strict = not allow_size_mismatch) | |
| if optimizer is not None: | |
| optimizer.load_state_dict(checkpoint['optimizer_state_dict']) | |
| return model, optimizer, checkpoint['epoch_idx'], checkpoint | |
| return model | |
| # -- logging utils | |
| def get_logger(args,save_path): | |
| log_path = '{}/{}_{}_{}classes_log.txt'.format(save_path,args.training_mode,args.lr,args.num_classes) | |
| logger = logging.getLogger("mylog") | |
| logger.setLevel(logging.INFO) | |
| fh = logging.FileHandler(log_path) | |
| fh.setLevel(logging.INFO) | |
| logger.addHandler(fh) | |
| console = logging.StreamHandler() | |
| console.setLevel(logging.INFO) | |
| logger.addHandler(console) | |
| return logger | |
| def update_logger_batch( args, logger, dset_loader, batch_idx, running_loss, running_corrects, running_all, batch_time, data_time ): | |
| perc_epoch = 100. * batch_idx / (len(dset_loader)-1) | |
| logger.info('[{:5.0f}/{:5.0f} ({:.0f}%)]\tLoss: {:.4f}\tAcc:{:.4f}\tCost time:{:1.3f} ({:1.3f})s\tData time:{:1.3f} ({:1.3f})\tInstances per second: {:.2f}'.format( | |
| running_all, | |
| len(dset_loader.dataset), | |
| perc_epoch, | |
| running_loss / running_all, | |
| running_corrects / running_all, | |
| batch_time.val, batch_time.avg, | |
| data_time.val, data_time.avg, | |
| args.batch_size/batch_time.avg )) | |
| def get_save_folder( args): | |
| # create save and log folder | |
| save_path = '{}/{}'.format( args.logging_dir, args.training_mode ) | |
| save_path += '/' + datetime.datetime.now().isoformat().split('.')[0] | |
| if not os.path.isdir(save_path): | |
| os.makedirs(save_path) | |
| return save_path | |