diff --git a/README.md b/README.md index ad35d6d..9638d09 100644 --- a/README.md +++ b/README.md @@ -147,5 +147,8 @@ If you find that this project helps your research, please consider citing the re If you want to contribute to this repo, please see [CONTRIBUTING.md](.github/CONTRIBUTING.md). Besides, please follow [CODE-OF-CONDUCT.md](.github/CODE-OF-CONDUCT.md). +We use `[black](https://github.com/psf/black)` for Python code formatter. +Please use `black . -l 120`. + # License The entire codebase is under the [MIT license](LICENSE.md). diff --git a/configs/qlib/workflow_config_mlp_Alpha360.yaml b/configs/qlib/workflow_config_mlp_Alpha360.yaml new file mode 100644 index 0000000..36582cf --- /dev/null +++ b/configs/qlib/workflow_config_mlp_Alpha360.yaml @@ -0,0 +1,82 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market all +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: DNNModelPytorch + module_path: qlib.contrib.model.pytorch_nn + kwargs: + loss: mse + input_dim: 360 + output_dim: 1 + lr: 0.002 + lr_decay: 0.96 + lr_decay_steps: 100 + optimizer: adam + max_steps: 8000 + batch_size: 4096 + GPU: 0 + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/configs/qlib/workflow_config_sfm_Alpha360.yaml b/configs/qlib/workflow_config_sfm_Alpha360.yaml new file mode 100644 index 0000000..8017e13 --- /dev/null +++ b/configs/qlib/workflow_config_sfm_Alpha360.yaml @@ -0,0 +1,85 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market all +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: SFM + module_path: qlib.contrib.model.pytorch_sfm + kwargs: + d_feat: 6 + hidden_size: 64 + output_dim: 32 + freq_dim: 25 + dropout_W: 0.5 + dropout_U: 0.5 + n_epochs: 20 + lr: 1e-3 + batch_size: 1600 + early_stop: 20 + eval_steps: 5 + loss: mse + optimizer: adam + GPU: 0 + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/exps/trading/baselines.py b/exps/trading/baselines.py index 00f3891..daccf43 100644 --- a/exps/trading/baselines.py +++ b/exps/trading/baselines.py @@ -4,6 +4,8 @@ # python exps/trading/baselines.py --alg GRU # python exps/trading/baselines.py --alg LSTM # python exps/trading/baselines.py --alg ALSTM +# python exps/trading/baselines.py --alg MLP +# python exps/trading/baselines.py --alg SFM # python exps/trading/baselines.py --alg XGBoost # python exps/trading/baselines.py --alg LightGBM ##################################################### @@ -17,6 +19,10 @@ lib_dir = (Path(__file__).parent / ".." / ".." / "lib").resolve() if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir)) +from procedures.q_exps import update_gpu +from procedures.q_exps import update_market +from procedures.q_exps import run_exp + import qlib from qlib.utils import init_instance_by_config from qlib.workflow import R @@ -31,15 +37,19 @@ def retrieve_configs(): alg2names = OrderedDict() alg2names["GRU"] = "workflow_config_gru_Alpha360.yaml" alg2names["LSTM"] = "workflow_config_lstm_Alpha360.yaml" + alg2names["MLP"] = "workflow_config_mlp_Alpha360.yaml" # A dual-stage attention-based recurrent neural network for time series prediction, IJCAI-2017 alg2names["ALSTM"] = "workflow_config_alstm_Alpha360.yaml" # XGBoost: A Scalable Tree Boosting System, KDD-2016 alg2names["XGBoost"] = "workflow_config_xgboost_Alpha360.yaml" # LightGBM: A Highly Efficient Gradient Boosting Decision Tree, NeurIPS-2017 alg2names["LightGBM"] = "workflow_config_lightgbm_Alpha360.yaml" + # State Frequency Memory (SFM): Stock Price Prediction via Discovering Multi-Frequency Trading Patterns, KDD-2017 + alg2names["SFM"] = "workflow_config_sfm_Alpha360.yaml" # find the yaml paths alg2paths = OrderedDict() + print("Start retrieving the algorithm configurations") for idx, (alg, name) in enumerate(alg2names.items()): path = config_dir / name assert path.exists(), "{:} does not exist.".format(path) @@ -48,56 +58,6 @@ def retrieve_configs(): return alg2paths -def update_gpu(config, gpu): - config = config.copy() - if "GPU" in config["task"]["model"]: - config["task"]["model"]["GPU"] = gpu - return config - - -def update_market(config, market): - config = config.copy() - config["market"] = market - config["data_handler_config"]["instruments"] = market - return config - - -def run_exp(task_config, dataset, experiment_name, recorder_name, uri): - - # model initiaiton - print("") - print("[{:}] - [{:}]: {:}".format(experiment_name, recorder_name, uri)) - print("dataset={:}".format(dataset)) - - model = init_instance_by_config(task_config["model"]) - - # start exp - with R.start(experiment_name=experiment_name, recorder_name=recorder_name, uri=uri): - - log_file = R.get_recorder().root_uri / "{:}.log".format(experiment_name) - set_log_basic_config(log_file) - - # train model - R.log_params(**flatten_dict(task_config)) - model.fit(dataset) - recorder = R.get_recorder() - R.save_objects(**{"model.pkl": model}) - - # generate records: prediction, backtest, and analysis - for record in task_config["record"]: - record = record.copy() - if record["class"] == "SignalRecord": - srconf = {"model": model, "dataset": dataset, "recorder": recorder} - record["kwargs"].update(srconf) - sr = init_instance_by_config(record) - sr.generate() - else: - rconf = {"recorder": recorder} - record["kwargs"].update(rconf) - ar = init_instance_by_config(record) - ar.generate() - - def main(xargs, exp_yaml): assert Path(exp_yaml).exists(), "{:} does not exist.".format(exp_yaml) diff --git a/lib/procedures/__init__.py b/lib/procedures/__init__.py index e9330c2..9392449 100644 --- a/lib/procedures/__init__.py +++ b/lib/procedures/__init__.py @@ -1,25 +1,36 @@ ################################################## # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 # ################################################## -from .starts import prepare_seed, prepare_logger, get_machine_info, save_checkpoint, copy_checkpoint +from .starts import prepare_seed +from .starts import prepare_logger +from .starts import get_machine_info +from .starts import save_checkpoint +from .starts import copy_checkpoint from .optimizers import get_optim_scheduler from .funcs_nasbench import evaluate_for_seed as bench_evaluate_for_seed from .funcs_nasbench import pure_evaluate as bench_pure_evaluate from .funcs_nasbench import get_nas_bench_loaders -def get_procedures(procedure): - from .basic_main import basic_train, basic_valid - from .search_main import search_train, search_valid - from .search_main_v2 import search_train_v2 - from .simple_KD_main import simple_KD_train, simple_KD_valid - train_funcs = {'basic' : basic_train, \ - 'search': search_train,'Simple-KD': simple_KD_train, \ - 'search-v2': search_train_v2} - valid_funcs = {'basic' : basic_valid, \ - 'search': search_valid,'Simple-KD': simple_KD_valid, \ - 'search-v2': search_valid} - - train_func = train_funcs[procedure] - valid_func = valid_funcs[procedure] - return train_func, valid_func +def get_procedures(procedure): + from .basic_main import basic_train, basic_valid + from .search_main import search_train, search_valid + from .search_main_v2 import search_train_v2 + from .simple_KD_main import simple_KD_train, simple_KD_valid + + train_funcs = { + "basic": basic_train, + "search": search_train, + "Simple-KD": simple_KD_train, + "search-v2": search_train_v2, + } + valid_funcs = { + "basic": basic_valid, + "search": search_valid, + "Simple-KD": simple_KD_valid, + "search-v2": search_valid, + } + + train_func = train_funcs[procedure] + valid_func = valid_funcs[procedure] + return train_func, valid_func diff --git a/lib/procedures/basic_main.py b/lib/procedures/basic_main.py index ed8cea4..3eba35b 100644 --- a/lib/procedures/basic_main.py +++ b/lib/procedures/basic_main.py @@ -3,73 +3,100 @@ ################################################## import os, sys, time, torch from log_utils import AverageMeter, time_string -from utils import obtain_accuracy +from utils import obtain_accuracy def basic_train(xloader, network, criterion, scheduler, optimizer, optim_config, extra_info, print_freq, logger): - loss, acc1, acc5 = procedure(xloader, network, criterion, scheduler, optimizer, 'train', optim_config, extra_info, print_freq, logger) - return loss, acc1, acc5 + loss, acc1, acc5 = procedure( + xloader, network, criterion, scheduler, optimizer, "train", optim_config, extra_info, print_freq, logger + ) + return loss, acc1, acc5 def basic_valid(xloader, network, criterion, optim_config, extra_info, print_freq, logger): - with torch.no_grad(): - loss, acc1, acc5 = procedure(xloader, network, criterion, None, None, 'valid', None, extra_info, print_freq, logger) - return loss, acc1, acc5 + with torch.no_grad(): + loss, acc1, acc5 = procedure( + xloader, network, criterion, None, None, "valid", None, extra_info, print_freq, logger + ) + return loss, acc1, acc5 def procedure(xloader, network, criterion, scheduler, optimizer, mode, config, extra_info, print_freq, logger): - data_time, batch_time, losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() - if mode == 'train': - network.train() - elif mode == 'valid': - network.eval() - else: raise ValueError("The mode is not right : {:}".format(mode)) - - #logger.log('[{:5s}] config :: auxiliary={:}, message={:}'.format(mode, config.auxiliary if hasattr(config, 'auxiliary') else -1, network.module.get_message())) - logger.log('[{:5s}] config :: auxiliary={:}'.format(mode, config.auxiliary if hasattr(config, 'auxiliary') else -1)) - end = time.time() - for i, (inputs, targets) in enumerate(xloader): - if mode == 'train': scheduler.update(None, 1.0 * i / len(xloader)) - # measure data loading time - data_time.update(time.time() - end) - # calculate prediction and loss - targets = targets.cuda(non_blocking=True) - - if mode == 'train': optimizer.zero_grad() - - features, logits = network(inputs) - if isinstance(logits, list): - assert len(logits) == 2, 'logits must has {:} items instead of {:}'.format(2, len(logits)) - logits, logits_aux = logits + data_time, batch_time, losses, top1, top5 = ( + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + ) + if mode == "train": + network.train() + elif mode == "valid": + network.eval() else: - logits, logits_aux = logits, None - loss = criterion(logits, targets) - if config is not None and hasattr(config, 'auxiliary') and config.auxiliary > 0: - loss_aux = criterion(logits_aux, targets) - loss += config.auxiliary * loss_aux - - if mode == 'train': - loss.backward() - optimizer.step() + raise ValueError("The mode is not right : {:}".format(mode)) - # record - prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - losses.update(loss.item(), inputs.size(0)) - top1.update (prec1.item(), inputs.size(0)) - top5.update (prec5.item(), inputs.size(0)) - - # measure elapsed time - batch_time.update(time.time() - end) + # logger.log('[{:5s}] config :: auxiliary={:}, message={:}'.format(mode, config.auxiliary if hasattr(config, 'auxiliary') else -1, network.module.get_message())) + logger.log( + "[{:5s}] config :: auxiliary={:}".format(mode, config.auxiliary if hasattr(config, "auxiliary") else -1) + ) end = time.time() + for i, (inputs, targets) in enumerate(xloader): + if mode == "train": + scheduler.update(None, 1.0 * i / len(xloader)) + # measure data loading time + data_time.update(time.time() - end) + # calculate prediction and loss + targets = targets.cuda(non_blocking=True) - if i % print_freq == 0 or (i+1) == len(xloader): - Sstr = ' {:5s} '.format(mode.upper()) + time_string() + ' [{:}][{:03d}/{:03d}]'.format(extra_info, i, len(xloader)) - if scheduler is not None: - Sstr += ' {:}'.format(scheduler.get_min_info()) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Lstr = 'Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})'.format(loss=losses, top1=top1, top5=top5) - Istr = 'Size={:}'.format(list(inputs.size())) - logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Istr) + if mode == "train": + optimizer.zero_grad() - logger.log(' **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}'.format(mode=mode.upper(), top1=top1, top5=top5, error1=100-top1.avg, error5=100-top5.avg, loss=losses.avg)) - return losses.avg, top1.avg, top5.avg + features, logits = network(inputs) + if isinstance(logits, list): + assert len(logits) == 2, "logits must has {:} items instead of {:}".format(2, len(logits)) + logits, logits_aux = logits + else: + logits, logits_aux = logits, None + loss = criterion(logits, targets) + if config is not None and hasattr(config, "auxiliary") and config.auxiliary > 0: + loss_aux = criterion(logits_aux, targets) + loss += config.auxiliary * loss_aux + + if mode == "train": + loss.backward() + optimizer.step() + + # record + prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) + losses.update(loss.item(), inputs.size(0)) + top1.update(prec1.item(), inputs.size(0)) + top5.update(prec5.item(), inputs.size(0)) + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + + if i % print_freq == 0 or (i + 1) == len(xloader): + Sstr = ( + " {:5s} ".format(mode.upper()) + + time_string() + + " [{:}][{:03d}/{:03d}]".format(extra_info, i, len(xloader)) + ) + if scheduler is not None: + Sstr += " {:}".format(scheduler.get_min_info()) + Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format( + batch_time=batch_time, data_time=data_time + ) + Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format( + loss=losses, top1=top1, top5=top5 + ) + Istr = "Size={:}".format(list(inputs.size())) + logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr) + + logger.log( + " **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}".format( + mode=mode.upper(), top1=top1, top5=top5, error1=100 - top1.avg, error5=100 - top5.avg, loss=losses.avg + ) + ) + return losses.avg, top1.avg, top5.avg diff --git a/lib/procedures/funcs_nasbench.py b/lib/procedures/funcs_nasbench.py index e1057a0..681d11d 100644 --- a/lib/procedures/funcs_nasbench.py +++ b/lib/procedures/funcs_nasbench.py @@ -5,199 +5,348 @@ import os, time, copy, torch, pathlib import datasets from config_utils import load_config -from procedures import prepare_seed, get_optim_scheduler -from utils import get_model_infos, obtain_accuracy -from log_utils import AverageMeter, time_string, convert_secs2time -from models import get_cell_based_tiny_net +from procedures import prepare_seed, get_optim_scheduler +from utils import get_model_infos, obtain_accuracy +from log_utils import AverageMeter, time_string, convert_secs2time +from models import get_cell_based_tiny_net -__all__ = ['evaluate_for_seed', 'pure_evaluate', 'get_nas_bench_loaders'] +__all__ = ["evaluate_for_seed", "pure_evaluate", "get_nas_bench_loaders"] def pure_evaluate(xloader, network, criterion=torch.nn.CrossEntropyLoss()): - data_time, batch_time, batch = AverageMeter(), AverageMeter(), None - losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter() - latencies, device = [], torch.cuda.current_device() - network.eval() - with torch.no_grad(): - end = time.time() - for i, (inputs, targets) in enumerate(xloader): - targets = targets.cuda(device=device, non_blocking=True) - inputs = inputs.cuda(device=device, non_blocking=True) - data_time.update(time.time() - end) - # forward - features, logits = network(inputs) - loss = criterion(logits, targets) - batch_time.update(time.time() - end) - if batch is None or batch == inputs.size(0): - batch = inputs.size(0) - latencies.append( batch_time.val - data_time.val ) - # record loss and accuracy - prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - losses.update(loss.item(), inputs.size(0)) - top1.update (prec1.item(), inputs.size(0)) - top5.update (prec5.item(), inputs.size(0)) - end = time.time() - if len(latencies) > 2: latencies = latencies[1:] - return losses.avg, top1.avg, top5.avg, latencies - + data_time, batch_time, batch = AverageMeter(), AverageMeter(), None + losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter() + latencies, device = [], torch.cuda.current_device() + network.eval() + with torch.no_grad(): + end = time.time() + for i, (inputs, targets) in enumerate(xloader): + targets = targets.cuda(device=device, non_blocking=True) + inputs = inputs.cuda(device=device, non_blocking=True) + data_time.update(time.time() - end) + # forward + features, logits = network(inputs) + loss = criterion(logits, targets) + batch_time.update(time.time() - end) + if batch is None or batch == inputs.size(0): + batch = inputs.size(0) + latencies.append(batch_time.val - data_time.val) + # record loss and accuracy + prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) + losses.update(loss.item(), inputs.size(0)) + top1.update(prec1.item(), inputs.size(0)) + top5.update(prec5.item(), inputs.size(0)) + end = time.time() + if len(latencies) > 2: + latencies = latencies[1:] + return losses.avg, top1.avg, top5.avg, latencies def procedure(xloader, network, criterion, scheduler, optimizer, mode: str): - losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter() - if mode == 'train' : network.train() - elif mode == 'valid': network.eval() - else: raise ValueError("The mode is not right : {:}".format(mode)) - device = torch.cuda.current_device() - data_time, batch_time, end = AverageMeter(), AverageMeter(), time.time() - for i, (inputs, targets) in enumerate(xloader): - if mode == 'train': scheduler.update(None, 1.0 * i / len(xloader)) + losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter() + if mode == "train": + network.train() + elif mode == "valid": + network.eval() + else: + raise ValueError("The mode is not right : {:}".format(mode)) + device = torch.cuda.current_device() + data_time, batch_time, end = AverageMeter(), AverageMeter(), time.time() + for i, (inputs, targets) in enumerate(xloader): + if mode == "train": + scheduler.update(None, 1.0 * i / len(xloader)) - targets = targets.cuda(device=device, non_blocking=True) - if mode == 'train': optimizer.zero_grad() - # forward - features, logits = network(inputs) - loss = criterion(logits, targets) - # backward - if mode == 'train': - loss.backward() - optimizer.step() - # record loss and accuracy - prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - losses.update(loss.item(), inputs.size(0)) - top1.update (prec1.item(), inputs.size(0)) - top5.update (prec5.item(), inputs.size(0)) - # count time - batch_time.update(time.time() - end) - end = time.time() - return losses.avg, top1.avg, top5.avg, batch_time.sum + targets = targets.cuda(device=device, non_blocking=True) + if mode == "train": + optimizer.zero_grad() + # forward + features, logits = network(inputs) + loss = criterion(logits, targets) + # backward + if mode == "train": + loss.backward() + optimizer.step() + # record loss and accuracy + prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) + losses.update(loss.item(), inputs.size(0)) + top1.update(prec1.item(), inputs.size(0)) + top5.update(prec5.item(), inputs.size(0)) + # count time + batch_time.update(time.time() - end) + end = time.time() + return losses.avg, top1.avg, top5.avg, batch_time.sum def evaluate_for_seed(arch_config, opt_config, train_loader, valid_loaders, seed: int, logger): - prepare_seed(seed) # random seed - net = get_cell_based_tiny_net(arch_config) - #net = TinyNetwork(arch_config['channel'], arch_config['num_cells'], arch, config.class_num) - flop, param = get_model_infos(net, opt_config.xshape) - logger.log('Network : {:}'.format(net.get_message()), False) - logger.log('{:} Seed-------------------------- {:} --------------------------'.format(time_string(), seed)) - logger.log('FLOP = {:} MB, Param = {:} MB'.format(flop, param)) - # train and valid - optimizer, scheduler, criterion = get_optim_scheduler(net.parameters(), opt_config) - default_device = torch.cuda.current_device() - network = torch.nn.DataParallel(net, device_ids=[default_device]).cuda(device=default_device) - criterion = criterion.cuda(device=default_device) - # start training - start_time, epoch_time, total_epoch = time.time(), AverageMeter(), opt_config.epochs + opt_config.warmup - train_losses, train_acc1es, train_acc5es, valid_losses, valid_acc1es, valid_acc5es = {}, {}, {}, {}, {}, {} - train_times , valid_times, lrs = {}, {}, {} - for epoch in range(total_epoch): - scheduler.update(epoch, 0.0) - lr = min(scheduler.get_lr()) - train_loss, train_acc1, train_acc5, train_tm = procedure(train_loader, network, criterion, scheduler, optimizer, 'train') - train_losses[epoch] = train_loss - train_acc1es[epoch] = train_acc1 - train_acc5es[epoch] = train_acc5 - train_times [epoch] = train_tm - lrs[epoch] = lr - with torch.no_grad(): - for key, xloder in valid_loaders.items(): - valid_loss, valid_acc1, valid_acc5, valid_tm = procedure(xloder , network, criterion, None, None, 'valid') - valid_losses['{:}@{:}'.format(key,epoch)] = valid_loss - valid_acc1es['{:}@{:}'.format(key,epoch)] = valid_acc1 - valid_acc5es['{:}@{:}'.format(key,epoch)] = valid_acc5 - valid_times ['{:}@{:}'.format(key,epoch)] = valid_tm + prepare_seed(seed) # random seed + net = get_cell_based_tiny_net(arch_config) + # net = TinyNetwork(arch_config['channel'], arch_config['num_cells'], arch, config.class_num) + flop, param = get_model_infos(net, opt_config.xshape) + logger.log("Network : {:}".format(net.get_message()), False) + logger.log("{:} Seed-------------------------- {:} --------------------------".format(time_string(), seed)) + logger.log("FLOP = {:} MB, Param = {:} MB".format(flop, param)) + # train and valid + optimizer, scheduler, criterion = get_optim_scheduler(net.parameters(), opt_config) + default_device = torch.cuda.current_device() + network = torch.nn.DataParallel(net, device_ids=[default_device]).cuda(device=default_device) + criterion = criterion.cuda(device=default_device) + # start training + start_time, epoch_time, total_epoch = time.time(), AverageMeter(), opt_config.epochs + opt_config.warmup + train_losses, train_acc1es, train_acc5es, valid_losses, valid_acc1es, valid_acc5es = {}, {}, {}, {}, {}, {} + train_times, valid_times, lrs = {}, {}, {} + for epoch in range(total_epoch): + scheduler.update(epoch, 0.0) + lr = min(scheduler.get_lr()) + train_loss, train_acc1, train_acc5, train_tm = procedure( + train_loader, network, criterion, scheduler, optimizer, "train" + ) + train_losses[epoch] = train_loss + train_acc1es[epoch] = train_acc1 + train_acc5es[epoch] = train_acc5 + train_times[epoch] = train_tm + lrs[epoch] = lr + with torch.no_grad(): + for key, xloder in valid_loaders.items(): + valid_loss, valid_acc1, valid_acc5, valid_tm = procedure( + xloder, network, criterion, None, None, "valid" + ) + valid_losses["{:}@{:}".format(key, epoch)] = valid_loss + valid_acc1es["{:}@{:}".format(key, epoch)] = valid_acc1 + valid_acc5es["{:}@{:}".format(key, epoch)] = valid_acc5 + valid_times["{:}@{:}".format(key, epoch)] = valid_tm - # measure elapsed time - epoch_time.update(time.time() - start_time) - start_time = time.time() - need_time = 'Time Left: {:}'.format( convert_secs2time(epoch_time.avg * (total_epoch-epoch-1), True) ) - logger.log('{:} {:} epoch={:03d}/{:03d} :: Train [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%] Valid [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%], lr={:}'.format(time_string(), need_time, epoch, total_epoch, train_loss, train_acc1, train_acc5, valid_loss, valid_acc1, valid_acc5, lr)) - info_seed = {'flop' : flop, - 'param': param, - 'arch_config' : arch_config._asdict(), - 'opt_config' : opt_config._asdict(), - 'total_epoch' : total_epoch , - 'train_losses': train_losses, - 'train_acc1es': train_acc1es, - 'train_acc5es': train_acc5es, - 'train_times' : train_times, - 'valid_losses': valid_losses, - 'valid_acc1es': valid_acc1es, - 'valid_acc5es': valid_acc5es, - 'valid_times' : valid_times, - 'learning_rates': lrs, - 'net_state_dict': net.state_dict(), - 'net_string' : '{:}'.format(net), - 'finish-train': True - } - return info_seed + # measure elapsed time + epoch_time.update(time.time() - start_time) + start_time = time.time() + need_time = "Time Left: {:}".format(convert_secs2time(epoch_time.avg * (total_epoch - epoch - 1), True)) + logger.log( + "{:} {:} epoch={:03d}/{:03d} :: Train [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%] Valid [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%], lr={:}".format( + time_string(), + need_time, + epoch, + total_epoch, + train_loss, + train_acc1, + train_acc5, + valid_loss, + valid_acc1, + valid_acc5, + lr, + ) + ) + info_seed = { + "flop": flop, + "param": param, + "arch_config": arch_config._asdict(), + "opt_config": opt_config._asdict(), + "total_epoch": total_epoch, + "train_losses": train_losses, + "train_acc1es": train_acc1es, + "train_acc5es": train_acc5es, + "train_times": train_times, + "valid_losses": valid_losses, + "valid_acc1es": valid_acc1es, + "valid_acc5es": valid_acc5es, + "valid_times": valid_times, + "learning_rates": lrs, + "net_state_dict": net.state_dict(), + "net_string": "{:}".format(net), + "finish-train": True, + } + return info_seed def get_nas_bench_loaders(workers): - torch.set_num_threads(workers) + torch.set_num_threads(workers) - root_dir = (pathlib.Path(__file__).parent / '..' / '..').resolve() - torch_dir = pathlib.Path(os.environ['TORCH_HOME']) - # cifar - cifar_config_path = root_dir / 'configs' / 'nas-benchmark' / 'CIFAR.config' - cifar_config = load_config(cifar_config_path, None, None) - get_datasets = datasets.get_datasets # a function to return the dataset - break_line = '-' * 150 - print ('{:} Create data-loader for all datasets'.format(time_string())) - print (break_line) - TRAIN_CIFAR10, VALID_CIFAR10, xshape, class_num = get_datasets('cifar10', str(torch_dir/'cifar.python'), -1) - print ('original CIFAR-10 : {:} training images and {:} test images : {:} input shape : {:} number of classes'.format(len(TRAIN_CIFAR10), len(VALID_CIFAR10), xshape, class_num)) - cifar10_splits = load_config(root_dir / 'configs' / 'nas-benchmark' / 'cifar-split.txt', None, None) - assert cifar10_splits.train[:10] == [0, 5, 7, 11, 13, 15, 16, 17, 20, 24] and cifar10_splits.valid[:10] == [1, 2, 3, 4, 6, 8, 9, 10, 12, 14] - temp_dataset = copy.deepcopy(TRAIN_CIFAR10) - temp_dataset.transform = VALID_CIFAR10.transform - # data loader - trainval_cifar10_loader = torch.utils.data.DataLoader(TRAIN_CIFAR10, batch_size=cifar_config.batch_size, shuffle=True , num_workers=workers, pin_memory=True) - train_cifar10_loader = torch.utils.data.DataLoader(TRAIN_CIFAR10, batch_size=cifar_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar10_splits.train), num_workers=workers, pin_memory=True) - valid_cifar10_loader = torch.utils.data.DataLoader(temp_dataset , batch_size=cifar_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar10_splits.valid), num_workers=workers, pin_memory=True) - test__cifar10_loader = torch.utils.data.DataLoader(VALID_CIFAR10, batch_size=cifar_config.batch_size, shuffle=False, num_workers=workers, pin_memory=True) - print ('CIFAR-10 : trval-loader has {:3d} batch with {:} per batch'.format(len(trainval_cifar10_loader), cifar_config.batch_size)) - print ('CIFAR-10 : train-loader has {:3d} batch with {:} per batch'.format(len(train_cifar10_loader), cifar_config.batch_size)) - print ('CIFAR-10 : valid-loader has {:3d} batch with {:} per batch'.format(len(valid_cifar10_loader), cifar_config.batch_size)) - print ('CIFAR-10 : test--loader has {:3d} batch with {:} per batch'.format(len(test__cifar10_loader), cifar_config.batch_size)) - print (break_line) - # CIFAR-100 - TRAIN_CIFAR100, VALID_CIFAR100, xshape, class_num = get_datasets('cifar100', str(torch_dir/'cifar.python'), -1) - print ('original CIFAR-100: {:} training images and {:} test images : {:} input shape : {:} number of classes'.format(len(TRAIN_CIFAR100), len(VALID_CIFAR100), xshape, class_num)) - cifar100_splits = load_config(root_dir / 'configs' / 'nas-benchmark' / 'cifar100-test-split.txt', None, None) - assert cifar100_splits.xvalid[:10] == [1, 3, 4, 5, 8, 10, 13, 14, 15, 16] and cifar100_splits.xtest[:10] == [0, 2, 6, 7, 9, 11, 12, 17, 20, 24] - train_cifar100_loader = torch.utils.data.DataLoader(TRAIN_CIFAR100, batch_size=cifar_config.batch_size, shuffle=True, num_workers=workers, pin_memory=True) - valid_cifar100_loader = torch.utils.data.DataLoader(VALID_CIFAR100, batch_size=cifar_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar100_splits.xvalid), num_workers=workers, pin_memory=True) - test__cifar100_loader = torch.utils.data.DataLoader(VALID_CIFAR100, batch_size=cifar_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar100_splits.xtest) , num_workers=workers, pin_memory=True) - print ('CIFAR-100 : train-loader has {:3d} batch'.format(len(train_cifar100_loader))) - print ('CIFAR-100 : valid-loader has {:3d} batch'.format(len(valid_cifar100_loader))) - print ('CIFAR-100 : test--loader has {:3d} batch'.format(len(test__cifar100_loader))) - print (break_line) + root_dir = (pathlib.Path(__file__).parent / ".." / "..").resolve() + torch_dir = pathlib.Path(os.environ["TORCH_HOME"]) + # cifar + cifar_config_path = root_dir / "configs" / "nas-benchmark" / "CIFAR.config" + cifar_config = load_config(cifar_config_path, None, None) + get_datasets = datasets.get_datasets # a function to return the dataset + break_line = "-" * 150 + print("{:} Create data-loader for all datasets".format(time_string())) + print(break_line) + TRAIN_CIFAR10, VALID_CIFAR10, xshape, class_num = get_datasets("cifar10", str(torch_dir / "cifar.python"), -1) + print( + "original CIFAR-10 : {:} training images and {:} test images : {:} input shape : {:} number of classes".format( + len(TRAIN_CIFAR10), len(VALID_CIFAR10), xshape, class_num + ) + ) + cifar10_splits = load_config(root_dir / "configs" / "nas-benchmark" / "cifar-split.txt", None, None) + assert cifar10_splits.train[:10] == [0, 5, 7, 11, 13, 15, 16, 17, 20, 24] and cifar10_splits.valid[:10] == [ + 1, + 2, + 3, + 4, + 6, + 8, + 9, + 10, + 12, + 14, + ] + temp_dataset = copy.deepcopy(TRAIN_CIFAR10) + temp_dataset.transform = VALID_CIFAR10.transform + # data loader + trainval_cifar10_loader = torch.utils.data.DataLoader( + TRAIN_CIFAR10, batch_size=cifar_config.batch_size, shuffle=True, num_workers=workers, pin_memory=True + ) + train_cifar10_loader = torch.utils.data.DataLoader( + TRAIN_CIFAR10, + batch_size=cifar_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar10_splits.train), + num_workers=workers, + pin_memory=True, + ) + valid_cifar10_loader = torch.utils.data.DataLoader( + temp_dataset, + batch_size=cifar_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar10_splits.valid), + num_workers=workers, + pin_memory=True, + ) + test__cifar10_loader = torch.utils.data.DataLoader( + VALID_CIFAR10, batch_size=cifar_config.batch_size, shuffle=False, num_workers=workers, pin_memory=True + ) + print( + "CIFAR-10 : trval-loader has {:3d} batch with {:} per batch".format( + len(trainval_cifar10_loader), cifar_config.batch_size + ) + ) + print( + "CIFAR-10 : train-loader has {:3d} batch with {:} per batch".format( + len(train_cifar10_loader), cifar_config.batch_size + ) + ) + print( + "CIFAR-10 : valid-loader has {:3d} batch with {:} per batch".format( + len(valid_cifar10_loader), cifar_config.batch_size + ) + ) + print( + "CIFAR-10 : test--loader has {:3d} batch with {:} per batch".format( + len(test__cifar10_loader), cifar_config.batch_size + ) + ) + print(break_line) + # CIFAR-100 + TRAIN_CIFAR100, VALID_CIFAR100, xshape, class_num = get_datasets("cifar100", str(torch_dir / "cifar.python"), -1) + print( + "original CIFAR-100: {:} training images and {:} test images : {:} input shape : {:} number of classes".format( + len(TRAIN_CIFAR100), len(VALID_CIFAR100), xshape, class_num + ) + ) + cifar100_splits = load_config(root_dir / "configs" / "nas-benchmark" / "cifar100-test-split.txt", None, None) + assert cifar100_splits.xvalid[:10] == [1, 3, 4, 5, 8, 10, 13, 14, 15, 16] and cifar100_splits.xtest[:10] == [ + 0, + 2, + 6, + 7, + 9, + 11, + 12, + 17, + 20, + 24, + ] + train_cifar100_loader = torch.utils.data.DataLoader( + TRAIN_CIFAR100, batch_size=cifar_config.batch_size, shuffle=True, num_workers=workers, pin_memory=True + ) + valid_cifar100_loader = torch.utils.data.DataLoader( + VALID_CIFAR100, + batch_size=cifar_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar100_splits.xvalid), + num_workers=workers, + pin_memory=True, + ) + test__cifar100_loader = torch.utils.data.DataLoader( + VALID_CIFAR100, + batch_size=cifar_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(cifar100_splits.xtest), + num_workers=workers, + pin_memory=True, + ) + print("CIFAR-100 : train-loader has {:3d} batch".format(len(train_cifar100_loader))) + print("CIFAR-100 : valid-loader has {:3d} batch".format(len(valid_cifar100_loader))) + print("CIFAR-100 : test--loader has {:3d} batch".format(len(test__cifar100_loader))) + print(break_line) - imagenet16_config_path = 'configs/nas-benchmark/ImageNet-16.config' - imagenet16_config = load_config(imagenet16_config_path, None, None) - TRAIN_ImageNet16_120, VALID_ImageNet16_120, xshape, class_num = get_datasets('ImageNet16-120', str(torch_dir/'cifar.python'/'ImageNet16'), -1) - print ('original TRAIN_ImageNet16_120: {:} training images and {:} test images : {:} input shape : {:} number of classes'.format(len(TRAIN_ImageNet16_120), len(VALID_ImageNet16_120), xshape, class_num)) - imagenet_splits = load_config(root_dir / 'configs' / 'nas-benchmark' / 'imagenet-16-120-test-split.txt', None, None) - assert imagenet_splits.xvalid[:10] == [1, 2, 3, 6, 7, 8, 9, 12, 16, 18] and imagenet_splits.xtest[:10] == [0, 4, 5, 10, 11, 13, 14, 15, 17, 20] - train_imagenet_loader = torch.utils.data.DataLoader(TRAIN_ImageNet16_120, batch_size=imagenet16_config.batch_size, shuffle=True, num_workers=workers, pin_memory=True) - valid_imagenet_loader = torch.utils.data.DataLoader(VALID_ImageNet16_120, batch_size=imagenet16_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(imagenet_splits.xvalid), num_workers=workers, pin_memory=True) - test__imagenet_loader = torch.utils.data.DataLoader(VALID_ImageNet16_120, batch_size=imagenet16_config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(imagenet_splits.xtest) , num_workers=workers, pin_memory=True) - print ('ImageNet-16-120 : train-loader has {:3d} batch with {:} per batch'.format(len(train_imagenet_loader), imagenet16_config.batch_size)) - print ('ImageNet-16-120 : valid-loader has {:3d} batch with {:} per batch'.format(len(valid_imagenet_loader), imagenet16_config.batch_size)) - print ('ImageNet-16-120 : test--loader has {:3d} batch with {:} per batch'.format(len(test__imagenet_loader), imagenet16_config.batch_size)) + imagenet16_config_path = "configs/nas-benchmark/ImageNet-16.config" + imagenet16_config = load_config(imagenet16_config_path, None, None) + TRAIN_ImageNet16_120, VALID_ImageNet16_120, xshape, class_num = get_datasets( + "ImageNet16-120", str(torch_dir / "cifar.python" / "ImageNet16"), -1 + ) + print( + "original TRAIN_ImageNet16_120: {:} training images and {:} test images : {:} input shape : {:} number of classes".format( + len(TRAIN_ImageNet16_120), len(VALID_ImageNet16_120), xshape, class_num + ) + ) + imagenet_splits = load_config(root_dir / "configs" / "nas-benchmark" / "imagenet-16-120-test-split.txt", None, None) + assert imagenet_splits.xvalid[:10] == [1, 2, 3, 6, 7, 8, 9, 12, 16, 18] and imagenet_splits.xtest[:10] == [ + 0, + 4, + 5, + 10, + 11, + 13, + 14, + 15, + 17, + 20, + ] + train_imagenet_loader = torch.utils.data.DataLoader( + TRAIN_ImageNet16_120, + batch_size=imagenet16_config.batch_size, + shuffle=True, + num_workers=workers, + pin_memory=True, + ) + valid_imagenet_loader = torch.utils.data.DataLoader( + VALID_ImageNet16_120, + batch_size=imagenet16_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(imagenet_splits.xvalid), + num_workers=workers, + pin_memory=True, + ) + test__imagenet_loader = torch.utils.data.DataLoader( + VALID_ImageNet16_120, + batch_size=imagenet16_config.batch_size, + sampler=torch.utils.data.sampler.SubsetRandomSampler(imagenet_splits.xtest), + num_workers=workers, + pin_memory=True, + ) + print( + "ImageNet-16-120 : train-loader has {:3d} batch with {:} per batch".format( + len(train_imagenet_loader), imagenet16_config.batch_size + ) + ) + print( + "ImageNet-16-120 : valid-loader has {:3d} batch with {:} per batch".format( + len(valid_imagenet_loader), imagenet16_config.batch_size + ) + ) + print( + "ImageNet-16-120 : test--loader has {:3d} batch with {:} per batch".format( + len(test__imagenet_loader), imagenet16_config.batch_size + ) + ) - # 'cifar10', 'cifar100', 'ImageNet16-120' - loaders = {'cifar10@trainval': trainval_cifar10_loader, - 'cifar10@train' : train_cifar10_loader, - 'cifar10@valid' : valid_cifar10_loader, - 'cifar10@test' : test__cifar10_loader, - 'cifar100@train' : train_cifar100_loader, - 'cifar100@valid' : valid_cifar100_loader, - 'cifar100@test' : test__cifar100_loader, - 'ImageNet16-120@train': train_imagenet_loader, - 'ImageNet16-120@valid': valid_imagenet_loader, - 'ImageNet16-120@test' : test__imagenet_loader} - return loaders \ No newline at end of file + # 'cifar10', 'cifar100', 'ImageNet16-120' + loaders = { + "cifar10@trainval": trainval_cifar10_loader, + "cifar10@train": train_cifar10_loader, + "cifar10@valid": valid_cifar10_loader, + "cifar10@test": test__cifar10_loader, + "cifar100@train": train_cifar100_loader, + "cifar100@valid": valid_cifar100_loader, + "cifar100@test": test__cifar100_loader, + "ImageNet16-120@train": train_imagenet_loader, + "ImageNet16-120@valid": valid_imagenet_loader, + "ImageNet16-120@test": test__imagenet_loader, + } + return loaders diff --git a/lib/procedures/optimizers.py b/lib/procedures/optimizers.py index 7fe086d..9f3143c 100644 --- a/lib/procedures/optimizers.py +++ b/lib/procedures/optimizers.py @@ -8,197 +8,201 @@ from torch.optim import Optimizer class _LRScheduler(object): + def __init__(self, optimizer, warmup_epochs, epochs): + if not isinstance(optimizer, Optimizer): + raise TypeError("{:} is not an Optimizer".format(type(optimizer).__name__)) + self.optimizer = optimizer + for group in optimizer.param_groups: + group.setdefault("initial_lr", group["lr"]) + self.base_lrs = list(map(lambda group: group["initial_lr"], optimizer.param_groups)) + self.max_epochs = epochs + self.warmup_epochs = warmup_epochs + self.current_epoch = 0 + self.current_iter = 0 - def __init__(self, optimizer, warmup_epochs, epochs): - if not isinstance(optimizer, Optimizer): - raise TypeError('{:} is not an Optimizer'.format(type(optimizer).__name__)) - self.optimizer = optimizer - for group in optimizer.param_groups: - group.setdefault('initial_lr', group['lr']) - self.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups)) - self.max_epochs = epochs - self.warmup_epochs = warmup_epochs - self.current_epoch = 0 - self.current_iter = 0 + def extra_repr(self): + return "" - def extra_repr(self): - return '' + def __repr__(self): + return "{name}(warmup={warmup_epochs}, max-epoch={max_epochs}, current::epoch={current_epoch}, iter={current_iter:.2f}".format( + name=self.__class__.__name__, **self.__dict__ + ) + ", {:})".format( + self.extra_repr() + ) - def __repr__(self): - return ('{name}(warmup={warmup_epochs}, max-epoch={max_epochs}, current::epoch={current_epoch}, iter={current_iter:.2f}'.format(name=self.__class__.__name__, **self.__dict__) - + ', {:})'.format(self.extra_repr())) + def state_dict(self): + return {key: value for key, value in self.__dict__.items() if key != "optimizer"} - def state_dict(self): - return {key: value for key, value in self.__dict__.items() if key != 'optimizer'} + def load_state_dict(self, state_dict): + self.__dict__.update(state_dict) - def load_state_dict(self, state_dict): - self.__dict__.update(state_dict) + def get_lr(self): + raise NotImplementedError - def get_lr(self): - raise NotImplementedError + def get_min_info(self): + lrs = self.get_lr() + return "#LR=[{:.6f}~{:.6f}] epoch={:03d}, iter={:4.2f}#".format( + min(lrs), max(lrs), self.current_epoch, self.current_iter + ) - def get_min_info(self): - lrs = self.get_lr() - return '#LR=[{:.6f}~{:.6f}] epoch={:03d}, iter={:4.2f}#'.format(min(lrs), max(lrs), self.current_epoch, self.current_iter) - - def get_min_lr(self): - return min( self.get_lr() ) - - def update(self, cur_epoch, cur_iter): - if cur_epoch is not None: - assert isinstance(cur_epoch, int) and cur_epoch>=0, 'invalid cur-epoch : {:}'.format(cur_epoch) - self.current_epoch = cur_epoch - if cur_iter is not None: - assert isinstance(cur_iter, float) and cur_iter>=0, 'invalid cur-iter : {:}'.format(cur_iter) - self.current_iter = cur_iter - for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): - param_group['lr'] = lr + def get_min_lr(self): + return min(self.get_lr()) + def update(self, cur_epoch, cur_iter): + if cur_epoch is not None: + assert isinstance(cur_epoch, int) and cur_epoch >= 0, "invalid cur-epoch : {:}".format(cur_epoch) + self.current_epoch = cur_epoch + if cur_iter is not None: + assert isinstance(cur_iter, float) and cur_iter >= 0, "invalid cur-iter : {:}".format(cur_iter) + self.current_iter = cur_iter + for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): + param_group["lr"] = lr class CosineAnnealingLR(_LRScheduler): + def __init__(self, optimizer, warmup_epochs, epochs, T_max, eta_min): + self.T_max = T_max + self.eta_min = eta_min + super(CosineAnnealingLR, self).__init__(optimizer, warmup_epochs, epochs) - def __init__(self, optimizer, warmup_epochs, epochs, T_max, eta_min): - self.T_max = T_max - self.eta_min = eta_min - super(CosineAnnealingLR, self).__init__(optimizer, warmup_epochs, epochs) - - def extra_repr(self): - return 'type={:}, T-max={:}, eta-min={:}'.format('cosine', self.T_max, self.eta_min) - - def get_lr(self): - lrs = [] - for base_lr in self.base_lrs: - if self.current_epoch >= self.warmup_epochs and self.current_epoch < self.max_epochs: - last_epoch = self.current_epoch - self.warmup_epochs - #if last_epoch < self.T_max: - #if last_epoch < self.max_epochs: - lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * last_epoch / self.T_max)) / 2 - #else: - # lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * (self.T_max-1.0) / self.T_max)) / 2 - elif self.current_epoch >= self.max_epochs: - lr = self.eta_min - else: - lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr - lrs.append( lr ) - return lrs + def extra_repr(self): + return "type={:}, T-max={:}, eta-min={:}".format("cosine", self.T_max, self.eta_min) + def get_lr(self): + lrs = [] + for base_lr in self.base_lrs: + if self.current_epoch >= self.warmup_epochs and self.current_epoch < self.max_epochs: + last_epoch = self.current_epoch - self.warmup_epochs + # if last_epoch < self.T_max: + # if last_epoch < self.max_epochs: + lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * last_epoch / self.T_max)) / 2 + # else: + # lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * (self.T_max-1.0) / self.T_max)) / 2 + elif self.current_epoch >= self.max_epochs: + lr = self.eta_min + else: + lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr + lrs.append(lr) + return lrs class MultiStepLR(_LRScheduler): + def __init__(self, optimizer, warmup_epochs, epochs, milestones, gammas): + assert len(milestones) == len(gammas), "invalid {:} vs {:}".format(len(milestones), len(gammas)) + self.milestones = milestones + self.gammas = gammas + super(MultiStepLR, self).__init__(optimizer, warmup_epochs, epochs) - def __init__(self, optimizer, warmup_epochs, epochs, milestones, gammas): - assert len(milestones) == len(gammas), 'invalid {:} vs {:}'.format(len(milestones), len(gammas)) - self.milestones = milestones - self.gammas = gammas - super(MultiStepLR, self).__init__(optimizer, warmup_epochs, epochs) + def extra_repr(self): + return "type={:}, milestones={:}, gammas={:}, base-lrs={:}".format( + "multistep", self.milestones, self.gammas, self.base_lrs + ) - def extra_repr(self): - return 'type={:}, milestones={:}, gammas={:}, base-lrs={:}'.format('multistep', self.milestones, self.gammas, self.base_lrs) - - def get_lr(self): - lrs = [] - for base_lr in self.base_lrs: - if self.current_epoch >= self.warmup_epochs: - last_epoch = self.current_epoch - self.warmup_epochs - idx = bisect_right(self.milestones, last_epoch) - lr = base_lr - for x in self.gammas[:idx]: lr *= x - else: - lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr - lrs.append( lr ) - return lrs + def get_lr(self): + lrs = [] + for base_lr in self.base_lrs: + if self.current_epoch >= self.warmup_epochs: + last_epoch = self.current_epoch - self.warmup_epochs + idx = bisect_right(self.milestones, last_epoch) + lr = base_lr + for x in self.gammas[:idx]: + lr *= x + else: + lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr + lrs.append(lr) + return lrs class ExponentialLR(_LRScheduler): + def __init__(self, optimizer, warmup_epochs, epochs, gamma): + self.gamma = gamma + super(ExponentialLR, self).__init__(optimizer, warmup_epochs, epochs) - def __init__(self, optimizer, warmup_epochs, epochs, gamma): - self.gamma = gamma - super(ExponentialLR, self).__init__(optimizer, warmup_epochs, epochs) + def extra_repr(self): + return "type={:}, gamma={:}, base-lrs={:}".format("exponential", self.gamma, self.base_lrs) - def extra_repr(self): - return 'type={:}, gamma={:}, base-lrs={:}'.format('exponential', self.gamma, self.base_lrs) - - def get_lr(self): - lrs = [] - for base_lr in self.base_lrs: - if self.current_epoch >= self.warmup_epochs: - last_epoch = self.current_epoch - self.warmup_epochs - assert last_epoch >= 0, 'invalid last_epoch : {:}'.format(last_epoch) - lr = base_lr * (self.gamma ** last_epoch) - else: - lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr - lrs.append( lr ) - return lrs + def get_lr(self): + lrs = [] + for base_lr in self.base_lrs: + if self.current_epoch >= self.warmup_epochs: + last_epoch = self.current_epoch - self.warmup_epochs + assert last_epoch >= 0, "invalid last_epoch : {:}".format(last_epoch) + lr = base_lr * (self.gamma ** last_epoch) + else: + lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr + lrs.append(lr) + return lrs class LinearLR(_LRScheduler): + def __init__(self, optimizer, warmup_epochs, epochs, max_LR, min_LR): + self.max_LR = max_LR + self.min_LR = min_LR + super(LinearLR, self).__init__(optimizer, warmup_epochs, epochs) - def __init__(self, optimizer, warmup_epochs, epochs, max_LR, min_LR): - self.max_LR = max_LR - self.min_LR = min_LR - super(LinearLR, self).__init__(optimizer, warmup_epochs, epochs) - - def extra_repr(self): - return 'type={:}, max_LR={:}, min_LR={:}, base-lrs={:}'.format('LinearLR', self.max_LR, self.min_LR, self.base_lrs) - - def get_lr(self): - lrs = [] - for base_lr in self.base_lrs: - if self.current_epoch >= self.warmup_epochs: - last_epoch = self.current_epoch - self.warmup_epochs - assert last_epoch >= 0, 'invalid last_epoch : {:}'.format(last_epoch) - ratio = (self.max_LR - self.min_LR) * last_epoch / self.max_epochs / self.max_LR - lr = base_lr * (1-ratio) - else: - lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr - lrs.append( lr ) - return lrs + def extra_repr(self): + return "type={:}, max_LR={:}, min_LR={:}, base-lrs={:}".format( + "LinearLR", self.max_LR, self.min_LR, self.base_lrs + ) + def get_lr(self): + lrs = [] + for base_lr in self.base_lrs: + if self.current_epoch >= self.warmup_epochs: + last_epoch = self.current_epoch - self.warmup_epochs + assert last_epoch >= 0, "invalid last_epoch : {:}".format(last_epoch) + ratio = (self.max_LR - self.min_LR) * last_epoch / self.max_epochs / self.max_LR + lr = base_lr * (1 - ratio) + else: + lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr + lrs.append(lr) + return lrs class CrossEntropyLabelSmooth(nn.Module): + def __init__(self, num_classes, epsilon): + super(CrossEntropyLabelSmooth, self).__init__() + self.num_classes = num_classes + self.epsilon = epsilon + self.logsoftmax = nn.LogSoftmax(dim=1) - def __init__(self, num_classes, epsilon): - super(CrossEntropyLabelSmooth, self).__init__() - self.num_classes = num_classes - self.epsilon = epsilon - self.logsoftmax = nn.LogSoftmax(dim=1) - - def forward(self, inputs, targets): - log_probs = self.logsoftmax(inputs) - targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) - targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes - loss = (-targets * log_probs).mean(0).sum() - return loss - + def forward(self, inputs, targets): + log_probs = self.logsoftmax(inputs) + targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) + targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes + loss = (-targets * log_probs).mean(0).sum() + return loss def get_optim_scheduler(parameters, config): - assert hasattr(config, 'optim') and hasattr(config, 'scheduler') and hasattr(config, 'criterion'), 'config must have optim / scheduler / criterion keys instead of {:}'.format(config) - if config.optim == 'SGD': - optim = torch.optim.SGD(parameters, config.LR, momentum=config.momentum, weight_decay=config.decay, nesterov=config.nesterov) - elif config.optim == 'RMSprop': - optim = torch.optim.RMSprop(parameters, config.LR, momentum=config.momentum, weight_decay=config.decay) - else: - raise ValueError('invalid optim : {:}'.format(config.optim)) + assert ( + hasattr(config, "optim") and hasattr(config, "scheduler") and hasattr(config, "criterion") + ), "config must have optim / scheduler / criterion keys instead of {:}".format(config) + if config.optim == "SGD": + optim = torch.optim.SGD( + parameters, config.LR, momentum=config.momentum, weight_decay=config.decay, nesterov=config.nesterov + ) + elif config.optim == "RMSprop": + optim = torch.optim.RMSprop(parameters, config.LR, momentum=config.momentum, weight_decay=config.decay) + else: + raise ValueError("invalid optim : {:}".format(config.optim)) - if config.scheduler == 'cos': - T_max = getattr(config, 'T_max', config.epochs) - scheduler = CosineAnnealingLR(optim, config.warmup, config.epochs, T_max, config.eta_min) - elif config.scheduler == 'multistep': - scheduler = MultiStepLR(optim, config.warmup, config.epochs, config.milestones, config.gammas) - elif config.scheduler == 'exponential': - scheduler = ExponentialLR(optim, config.warmup, config.epochs, config.gamma) - elif config.scheduler == 'linear': - scheduler = LinearLR(optim, config.warmup, config.epochs, config.LR, config.LR_min) - else: - raise ValueError('invalid scheduler : {:}'.format(config.scheduler)) + if config.scheduler == "cos": + T_max = getattr(config, "T_max", config.epochs) + scheduler = CosineAnnealingLR(optim, config.warmup, config.epochs, T_max, config.eta_min) + elif config.scheduler == "multistep": + scheduler = MultiStepLR(optim, config.warmup, config.epochs, config.milestones, config.gammas) + elif config.scheduler == "exponential": + scheduler = ExponentialLR(optim, config.warmup, config.epochs, config.gamma) + elif config.scheduler == "linear": + scheduler = LinearLR(optim, config.warmup, config.epochs, config.LR, config.LR_min) + else: + raise ValueError("invalid scheduler : {:}".format(config.scheduler)) - if config.criterion == 'Softmax': - criterion = torch.nn.CrossEntropyLoss() - elif config.criterion == 'SmoothSoftmax': - criterion = CrossEntropyLabelSmooth(config.class_num, config.label_smooth) - else: - raise ValueError('invalid criterion : {:}'.format(config.criterion)) - return optim, scheduler, criterion + if config.criterion == "Softmax": + criterion = torch.nn.CrossEntropyLoss() + elif config.criterion == "SmoothSoftmax": + criterion = CrossEntropyLabelSmooth(config.class_num, config.label_smooth) + else: + raise ValueError("invalid criterion : {:}".format(config.criterion)) + return optim, scheduler, criterion diff --git a/lib/procedures/q_exps.py b/lib/procedures/q_exps.py index a19d6a8..3b62167 100644 --- a/lib/procedures/q_exps.py +++ b/lib/procedures/q_exps.py @@ -7,11 +7,12 @@ from qlib.utils import init_instance_by_config from qlib.workflow import R from qlib.utils import flatten_dict from qlib.log import set_log_basic_config +from qlib.log import get_module_logger def update_gpu(config, gpu): config = config.copy() - if "task" in config and "GPU" in config["task"]["model"]: + if "task" in config and "moodel" in config["task"] and "GPU" in config["task"]["model"]: config["task"]["model"]["GPU"] = gpu elif "model" in config and "GPU" in config["model"]: config["model"]["GPU"] = gpu @@ -29,11 +30,6 @@ def update_market(config, market): def run_exp(task_config, dataset, experiment_name, recorder_name, uri): - # model initiaiton - print("") - print("[{:}] - [{:}]: {:}".format(experiment_name, recorder_name, uri)) - print("dataset={:}".format(dataset)) - model = init_instance_by_config(task_config["model"]) # start exp @@ -41,6 +37,10 @@ def run_exp(task_config, dataset, experiment_name, recorder_name, uri): log_file = R.get_recorder().root_uri / "{:}.log".format(experiment_name) set_log_basic_config(log_file) + logger = get_module_logger("q.run_exp") + logger.info("task_config={:}".format(task_config)) + logger.info("[{:}] - [{:}]: {:}".format(experiment_name, recorder_name, uri)) + logger.info("dataset={:}".format(dataset)) # train model R.log_params(**flatten_dict(task_config)) diff --git a/lib/procedures/search_main.py b/lib/procedures/search_main.py index 48ed44f..ecdea00 100644 --- a/lib/procedures/search_main.py +++ b/lib/procedures/search_main.py @@ -3,124 +3,170 @@ ################################################## import os, sys, time, torch from log_utils import AverageMeter, time_string -from utils import obtain_accuracy -from models import change_key +from utils import obtain_accuracy +from models import change_key def get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant): - expected_flop = torch.mean( expected_flop ) + expected_flop = torch.mean(expected_flop) - if flop_cur < flop_need - flop_tolerant: # Too Small FLOP - loss = - torch.log( expected_flop ) - #elif flop_cur > flop_need + flop_tolerant: # Too Large FLOP - elif flop_cur > flop_need: # Too Large FLOP - loss = torch.log( expected_flop ) - else: # Required FLOP - loss = None - if loss is None: return 0, 0 - else : return loss, loss.item() + if flop_cur < flop_need - flop_tolerant: # Too Small FLOP + loss = -torch.log(expected_flop) + # elif flop_cur > flop_need + flop_tolerant: # Too Large FLOP + elif flop_cur > flop_need: # Too Large FLOP + loss = torch.log(expected_flop) + else: # Required FLOP + loss = None + if loss is None: + return 0, 0 + else: + return loss, loss.item() -def search_train(search_loader, network, criterion, scheduler, base_optimizer, arch_optimizer, optim_config, extra_info, print_freq, logger): - data_time, batch_time = AverageMeter(), AverageMeter() - base_losses, arch_losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() - arch_cls_losses, arch_flop_losses = AverageMeter(), AverageMeter() - epoch_str, flop_need, flop_weight, flop_tolerant = extra_info['epoch-str'], extra_info['FLOP-exp'], extra_info['FLOP-weight'], extra_info['FLOP-tolerant'] +def search_train( + search_loader, + network, + criterion, + scheduler, + base_optimizer, + arch_optimizer, + optim_config, + extra_info, + print_freq, + logger, +): + data_time, batch_time = AverageMeter(), AverageMeter() + base_losses, arch_losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() + arch_cls_losses, arch_flop_losses = AverageMeter(), AverageMeter() + epoch_str, flop_need, flop_weight, flop_tolerant = ( + extra_info["epoch-str"], + extra_info["FLOP-exp"], + extra_info["FLOP-weight"], + extra_info["FLOP-tolerant"], + ) - network.train() - logger.log('[Search] : {:}, FLOP-Require={:.2f} MB, FLOP-WEIGHT={:.2f}'.format(epoch_str, flop_need, flop_weight)) - end = time.time() - network.apply( change_key('search_mode', 'search') ) - for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(search_loader): - scheduler.update(None, 1.0 * step / len(search_loader)) - # calculate prediction and loss - base_targets = base_targets.cuda(non_blocking=True) - arch_targets = arch_targets.cuda(non_blocking=True) - # measure data loading time - data_time.update(time.time() - end) - - # update the weights - base_optimizer.zero_grad() - logits, expected_flop = network(base_inputs) - #network.apply( change_key('search_mode', 'basic') ) - #features, logits = network(base_inputs) - base_loss = criterion(logits, base_targets) - base_loss.backward() - base_optimizer.step() - # record - prec1, prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5)) - base_losses.update(base_loss.item(), base_inputs.size(0)) - top1.update (prec1.item(), base_inputs.size(0)) - top5.update (prec5.item(), base_inputs.size(0)) - - # update the architecture - arch_optimizer.zero_grad() - logits, expected_flop = network(arch_inputs) - flop_cur = network.module.get_flop('genotype', None, None) - flop_loss, flop_loss_scale = get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant) - acls_loss = criterion(logits, arch_targets) - arch_loss = acls_loss + flop_loss * flop_weight - arch_loss.backward() - arch_optimizer.step() - - # record - arch_losses.update(arch_loss.item(), arch_inputs.size(0)) - arch_flop_losses.update(flop_loss_scale, arch_inputs.size(0)) - arch_cls_losses.update (acls_loss.item(), arch_inputs.size(0)) - - # measure elapsed time - batch_time.update(time.time() - end) + network.train() + logger.log("[Search] : {:}, FLOP-Require={:.2f} MB, FLOP-WEIGHT={:.2f}".format(epoch_str, flop_need, flop_weight)) end = time.time() - if step % print_freq == 0 or (step+1) == len(search_loader): - Sstr = '**TRAIN** ' + time_string() + ' [{:}][{:03d}/{:03d}]'.format(epoch_str, step, len(search_loader)) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Lstr = 'Base-Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})'.format(loss=base_losses, top1=top1, top5=top5) - Vstr = 'Acls-loss {aloss.val:.3f} ({aloss.avg:.3f}) FLOP-Loss {floss.val:.3f} ({floss.avg:.3f}) Arch-Loss {loss.val:.3f} ({loss.avg:.3f})'.format(aloss=arch_cls_losses, floss=arch_flop_losses, loss=arch_losses) - logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr) - #Istr = 'Bsz={:} Asz={:}'.format(list(base_inputs.size()), list(arch_inputs.size())) - #logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' ' + Istr) - #print(network.module.get_arch_info()) - #print(network.module.width_attentions[0]) - #print(network.module.width_attentions[1]) + network.apply(change_key("search_mode", "search")) + for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(search_loader): + scheduler.update(None, 1.0 * step / len(search_loader)) + # calculate prediction and loss + base_targets = base_targets.cuda(non_blocking=True) + arch_targets = arch_targets.cuda(non_blocking=True) + # measure data loading time + data_time.update(time.time() - end) - logger.log(' **TRAIN** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Base-Loss:{baseloss:.3f}, Arch-Loss={archloss:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg, error5=100-top5.avg, baseloss=base_losses.avg, archloss=arch_losses.avg)) - return base_losses.avg, arch_losses.avg, top1.avg, top5.avg + # update the weights + base_optimizer.zero_grad() + logits, expected_flop = network(base_inputs) + # network.apply( change_key('search_mode', 'basic') ) + # features, logits = network(base_inputs) + base_loss = criterion(logits, base_targets) + base_loss.backward() + base_optimizer.step() + # record + prec1, prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5)) + base_losses.update(base_loss.item(), base_inputs.size(0)) + top1.update(prec1.item(), base_inputs.size(0)) + top5.update(prec5.item(), base_inputs.size(0)) + # update the architecture + arch_optimizer.zero_grad() + logits, expected_flop = network(arch_inputs) + flop_cur = network.module.get_flop("genotype", None, None) + flop_loss, flop_loss_scale = get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant) + acls_loss = criterion(logits, arch_targets) + arch_loss = acls_loss + flop_loss * flop_weight + arch_loss.backward() + arch_optimizer.step() + + # record + arch_losses.update(arch_loss.item(), arch_inputs.size(0)) + arch_flop_losses.update(flop_loss_scale, arch_inputs.size(0)) + arch_cls_losses.update(acls_loss.item(), arch_inputs.size(0)) + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + if step % print_freq == 0 or (step + 1) == len(search_loader): + Sstr = "**TRAIN** " + time_string() + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(search_loader)) + Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format( + batch_time=batch_time, data_time=data_time + ) + Lstr = "Base-Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format( + loss=base_losses, top1=top1, top5=top5 + ) + Vstr = "Acls-loss {aloss.val:.3f} ({aloss.avg:.3f}) FLOP-Loss {floss.val:.3f} ({floss.avg:.3f}) Arch-Loss {loss.val:.3f} ({loss.avg:.3f})".format( + aloss=arch_cls_losses, floss=arch_flop_losses, loss=arch_losses + ) + logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Vstr) + # Istr = 'Bsz={:} Asz={:}'.format(list(base_inputs.size()), list(arch_inputs.size())) + # logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' ' + Istr) + # print(network.module.get_arch_info()) + # print(network.module.width_attentions[0]) + # print(network.module.width_attentions[1]) + + logger.log( + " **TRAIN** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Base-Loss:{baseloss:.3f}, Arch-Loss={archloss:.3f}".format( + top1=top1, + top5=top5, + error1=100 - top1.avg, + error5=100 - top5.avg, + baseloss=base_losses.avg, + archloss=arch_losses.avg, + ) + ) + return base_losses.avg, arch_losses.avg, top1.avg, top5.avg def search_valid(xloader, network, criterion, extra_info, print_freq, logger): - data_time, batch_time, losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() + data_time, batch_time, losses, top1, top5 = ( + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + ) - network.eval() - network.apply( change_key('search_mode', 'search') ) - end = time.time() - #logger.log('Starting evaluating {:}'.format(epoch_info)) - with torch.no_grad(): - for i, (inputs, targets) in enumerate(xloader): - # measure data loading time - data_time.update(time.time() - end) - # calculate prediction and loss - targets = targets.cuda(non_blocking=True) + network.eval() + network.apply(change_key("search_mode", "search")) + end = time.time() + # logger.log('Starting evaluating {:}'.format(epoch_info)) + with torch.no_grad(): + for i, (inputs, targets) in enumerate(xloader): + # measure data loading time + data_time.update(time.time() - end) + # calculate prediction and loss + targets = targets.cuda(non_blocking=True) - logits, expected_flop = network(inputs) - loss = criterion(logits, targets) - # record - prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - losses.update(loss.item(), inputs.size(0)) - top1.update (prec1.item(), inputs.size(0)) - top5.update (prec5.item(), inputs.size(0)) + logits, expected_flop = network(inputs) + loss = criterion(logits, targets) + # record + prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) + losses.update(loss.item(), inputs.size(0)) + top1.update(prec1.item(), inputs.size(0)) + top5.update(prec5.item(), inputs.size(0)) - # measure elapsed time - batch_time.update(time.time() - end) - end = time.time() + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() - if i % print_freq == 0 or (i+1) == len(xloader): - Sstr = '**VALID** ' + time_string() + ' [{:}][{:03d}/{:03d}]'.format(extra_info, i, len(xloader)) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Lstr = 'Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})'.format(loss=losses, top1=top1, top5=top5) - Istr = 'Size={:}'.format(list(inputs.size())) - logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Istr) + if i % print_freq == 0 or (i + 1) == len(xloader): + Sstr = "**VALID** " + time_string() + " [{:}][{:03d}/{:03d}]".format(extra_info, i, len(xloader)) + Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format( + batch_time=batch_time, data_time=data_time + ) + Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format( + loss=losses, top1=top1, top5=top5 + ) + Istr = "Size={:}".format(list(inputs.size())) + logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr) - logger.log(' **VALID** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg, error5=100-top5.avg, loss=losses.avg)) - - return losses.avg, top1.avg, top5.avg + logger.log( + " **VALID** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}".format( + top1=top1, top5=top5, error1=100 - top1.avg, error5=100 - top5.avg, loss=losses.avg + ) + ) + + return losses.avg, top1.avg, top5.avg diff --git a/lib/procedures/search_main_v2.py b/lib/procedures/search_main_v2.py index 46707ef..0b1fbca 100644 --- a/lib/procedures/search_main_v2.py +++ b/lib/procedures/search_main_v2.py @@ -3,85 +3,118 @@ ################################################## import os, sys, time, torch from log_utils import AverageMeter, time_string -from utils import obtain_accuracy -from models import change_key +from utils import obtain_accuracy +from models import change_key def get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant): - expected_flop = torch.mean( expected_flop ) + expected_flop = torch.mean(expected_flop) - if flop_cur < flop_need - flop_tolerant: # Too Small FLOP - loss = - torch.log( expected_flop ) - #elif flop_cur > flop_need + flop_tolerant: # Too Large FLOP - elif flop_cur > flop_need: # Too Large FLOP - loss = torch.log( expected_flop ) - else: # Required FLOP - loss = None - if loss is None: return 0, 0 - else : return loss, loss.item() + if flop_cur < flop_need - flop_tolerant: # Too Small FLOP + loss = -torch.log(expected_flop) + # elif flop_cur > flop_need + flop_tolerant: # Too Large FLOP + elif flop_cur > flop_need: # Too Large FLOP + loss = torch.log(expected_flop) + else: # Required FLOP + loss = None + if loss is None: + return 0, 0 + else: + return loss, loss.item() -def search_train_v2(search_loader, network, criterion, scheduler, base_optimizer, arch_optimizer, optim_config, extra_info, print_freq, logger): - data_time, batch_time = AverageMeter(), AverageMeter() - base_losses, arch_losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() - arch_cls_losses, arch_flop_losses = AverageMeter(), AverageMeter() - epoch_str, flop_need, flop_weight, flop_tolerant = extra_info['epoch-str'], extra_info['FLOP-exp'], extra_info['FLOP-weight'], extra_info['FLOP-tolerant'] +def search_train_v2( + search_loader, + network, + criterion, + scheduler, + base_optimizer, + arch_optimizer, + optim_config, + extra_info, + print_freq, + logger, +): + data_time, batch_time = AverageMeter(), AverageMeter() + base_losses, arch_losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() + arch_cls_losses, arch_flop_losses = AverageMeter(), AverageMeter() + epoch_str, flop_need, flop_weight, flop_tolerant = ( + extra_info["epoch-str"], + extra_info["FLOP-exp"], + extra_info["FLOP-weight"], + extra_info["FLOP-tolerant"], + ) - network.train() - logger.log('[Search] : {:}, FLOP-Require={:.2f} MB, FLOP-WEIGHT={:.2f}'.format(epoch_str, flop_need, flop_weight)) - end = time.time() - network.apply( change_key('search_mode', 'search') ) - for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(search_loader): - scheduler.update(None, 1.0 * step / len(search_loader)) - # calculate prediction and loss - base_targets = base_targets.cuda(non_blocking=True) - arch_targets = arch_targets.cuda(non_blocking=True) - # measure data loading time - data_time.update(time.time() - end) - - # update the weights - base_optimizer.zero_grad() - logits, expected_flop = network(base_inputs) - base_loss = criterion(logits, base_targets) - base_loss.backward() - base_optimizer.step() - # record - prec1, prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5)) - base_losses.update(base_loss.item(), base_inputs.size(0)) - top1.update (prec1.item(), base_inputs.size(0)) - top5.update (prec5.item(), base_inputs.size(0)) - - # update the architecture - arch_optimizer.zero_grad() - logits, expected_flop = network(arch_inputs) - flop_cur = network.module.get_flop('genotype', None, None) - flop_loss, flop_loss_scale = get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant) - acls_loss = criterion(logits, arch_targets) - arch_loss = acls_loss + flop_loss * flop_weight - arch_loss.backward() - arch_optimizer.step() - - # record - arch_losses.update(arch_loss.item(), arch_inputs.size(0)) - arch_flop_losses.update(flop_loss_scale, arch_inputs.size(0)) - arch_cls_losses.update (acls_loss.item(), arch_inputs.size(0)) - - # measure elapsed time - batch_time.update(time.time() - end) + network.train() + logger.log("[Search] : {:}, FLOP-Require={:.2f} MB, FLOP-WEIGHT={:.2f}".format(epoch_str, flop_need, flop_weight)) end = time.time() - if step % print_freq == 0 or (step+1) == len(search_loader): - Sstr = '**TRAIN** ' + time_string() + ' [{:}][{:03d}/{:03d}]'.format(epoch_str, step, len(search_loader)) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Lstr = 'Base-Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})'.format(loss=base_losses, top1=top1, top5=top5) - Vstr = 'Acls-loss {aloss.val:.3f} ({aloss.avg:.3f}) FLOP-Loss {floss.val:.3f} ({floss.avg:.3f}) Arch-Loss {loss.val:.3f} ({loss.avg:.3f})'.format(aloss=arch_cls_losses, floss=arch_flop_losses, loss=arch_losses) - logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr) - #num_bytes = torch.cuda.max_memory_allocated( next(network.parameters()).device ) * 1.0 - #logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' GPU={:.2f}MB'.format(num_bytes/1e6)) - #Istr = 'Bsz={:} Asz={:}'.format(list(base_inputs.size()), list(arch_inputs.size())) - #logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' ' + Istr) - #print(network.module.get_arch_info()) - #print(network.module.width_attentions[0]) - #print(network.module.width_attentions[1]) + network.apply(change_key("search_mode", "search")) + for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(search_loader): + scheduler.update(None, 1.0 * step / len(search_loader)) + # calculate prediction and loss + base_targets = base_targets.cuda(non_blocking=True) + arch_targets = arch_targets.cuda(non_blocking=True) + # measure data loading time + data_time.update(time.time() - end) - logger.log(' **TRAIN** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Base-Loss:{baseloss:.3f}, Arch-Loss={archloss:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg, error5=100-top5.avg, baseloss=base_losses.avg, archloss=arch_losses.avg)) - return base_losses.avg, arch_losses.avg, top1.avg, top5.avg + # update the weights + base_optimizer.zero_grad() + logits, expected_flop = network(base_inputs) + base_loss = criterion(logits, base_targets) + base_loss.backward() + base_optimizer.step() + # record + prec1, prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5)) + base_losses.update(base_loss.item(), base_inputs.size(0)) + top1.update(prec1.item(), base_inputs.size(0)) + top5.update(prec5.item(), base_inputs.size(0)) + + # update the architecture + arch_optimizer.zero_grad() + logits, expected_flop = network(arch_inputs) + flop_cur = network.module.get_flop("genotype", None, None) + flop_loss, flop_loss_scale = get_flop_loss(expected_flop, flop_cur, flop_need, flop_tolerant) + acls_loss = criterion(logits, arch_targets) + arch_loss = acls_loss + flop_loss * flop_weight + arch_loss.backward() + arch_optimizer.step() + + # record + arch_losses.update(arch_loss.item(), arch_inputs.size(0)) + arch_flop_losses.update(flop_loss_scale, arch_inputs.size(0)) + arch_cls_losses.update(acls_loss.item(), arch_inputs.size(0)) + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + if step % print_freq == 0 or (step + 1) == len(search_loader): + Sstr = "**TRAIN** " + time_string() + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(search_loader)) + Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format( + batch_time=batch_time, data_time=data_time + ) + Lstr = "Base-Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format( + loss=base_losses, top1=top1, top5=top5 + ) + Vstr = "Acls-loss {aloss.val:.3f} ({aloss.avg:.3f}) FLOP-Loss {floss.val:.3f} ({floss.avg:.3f}) Arch-Loss {loss.val:.3f} ({loss.avg:.3f})".format( + aloss=arch_cls_losses, floss=arch_flop_losses, loss=arch_losses + ) + logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Vstr) + # num_bytes = torch.cuda.max_memory_allocated( next(network.parameters()).device ) * 1.0 + # logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' GPU={:.2f}MB'.format(num_bytes/1e6)) + # Istr = 'Bsz={:} Asz={:}'.format(list(base_inputs.size()), list(arch_inputs.size())) + # logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' ' + Istr) + # print(network.module.get_arch_info()) + # print(network.module.width_attentions[0]) + # print(network.module.width_attentions[1]) + + logger.log( + " **TRAIN** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Base-Loss:{baseloss:.3f}, Arch-Loss={archloss:.3f}".format( + top1=top1, + top5=top5, + error1=100 - top1.avg, + error5=100 - top5.avg, + baseloss=base_losses.avg, + archloss=arch_losses.avg, + ) + ) + return base_losses.avg, arch_losses.avg, top1.avg, top5.avg diff --git a/lib/procedures/simple_KD_main.py b/lib/procedures/simple_KD_main.py index 5ad8c3b..cc7b3da 100644 --- a/lib/procedures/simple_KD_main.py +++ b/lib/procedures/simple_KD_main.py @@ -3,92 +3,143 @@ ##################################################### import os, sys, time, torch import torch.nn.functional as F + # our modules from log_utils import AverageMeter, time_string -from utils import obtain_accuracy +from utils import obtain_accuracy -def simple_KD_train(xloader, teacher, network, criterion, scheduler, optimizer, optim_config, extra_info, print_freq, logger): - loss, acc1, acc5 = procedure(xloader, teacher, network, criterion, scheduler, optimizer, 'train', optim_config, extra_info, print_freq, logger) - return loss, acc1, acc5 +def simple_KD_train( + xloader, teacher, network, criterion, scheduler, optimizer, optim_config, extra_info, print_freq, logger +): + loss, acc1, acc5 = procedure( + xloader, + teacher, + network, + criterion, + scheduler, + optimizer, + "train", + optim_config, + extra_info, + print_freq, + logger, + ) + return loss, acc1, acc5 + def simple_KD_valid(xloader, teacher, network, criterion, optim_config, extra_info, print_freq, logger): - with torch.no_grad(): - loss, acc1, acc5 = procedure(xloader, teacher, network, criterion, None, None, 'valid', optim_config, extra_info, print_freq, logger) - return loss, acc1, acc5 + with torch.no_grad(): + loss, acc1, acc5 = procedure( + xloader, teacher, network, criterion, None, None, "valid", optim_config, extra_info, print_freq, logger + ) + return loss, acc1, acc5 -def loss_KD_fn(criterion, student_logits, teacher_logits, studentFeatures, teacherFeatures, targets, alpha, temperature): - basic_loss = criterion(student_logits, targets) * (1. - alpha) - log_student= F.log_softmax(student_logits / temperature, dim=1) - sof_teacher= F.softmax (teacher_logits / temperature, dim=1) - KD_loss = F.kl_div(log_student, sof_teacher, reduction='batchmean') * (alpha * temperature * temperature) - return basic_loss + KD_loss +def loss_KD_fn( + criterion, student_logits, teacher_logits, studentFeatures, teacherFeatures, targets, alpha, temperature +): + basic_loss = criterion(student_logits, targets) * (1.0 - alpha) + log_student = F.log_softmax(student_logits / temperature, dim=1) + sof_teacher = F.softmax(teacher_logits / temperature, dim=1) + KD_loss = F.kl_div(log_student, sof_teacher, reduction="batchmean") * (alpha * temperature * temperature) + return basic_loss + KD_loss def procedure(xloader, teacher, network, criterion, scheduler, optimizer, mode, config, extra_info, print_freq, logger): - data_time, batch_time, losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() - Ttop1, Ttop5 = AverageMeter(), AverageMeter() - if mode == 'train': - network.train() - elif mode == 'valid': - network.eval() - else: raise ValueError("The mode is not right : {:}".format(mode)) - teacher.eval() - - logger.log('[{:5s}] config :: auxiliary={:}, KD :: [alpha={:.2f}, temperature={:.2f}]'.format(mode, config.auxiliary if hasattr(config, 'auxiliary') else -1, config.KD_alpha, config.KD_temperature)) - end = time.time() - for i, (inputs, targets) in enumerate(xloader): - if mode == 'train': scheduler.update(None, 1.0 * i / len(xloader)) - # measure data loading time - data_time.update(time.time() - end) - # calculate prediction and loss - targets = targets.cuda(non_blocking=True) - - if mode == 'train': optimizer.zero_grad() - - student_f, logits = network(inputs) - if isinstance(logits, list): - assert len(logits) == 2, 'logits must has {:} items instead of {:}'.format(2, len(logits)) - logits, logits_aux = logits + data_time, batch_time, losses, top1, top5 = ( + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + AverageMeter(), + ) + Ttop1, Ttop5 = AverageMeter(), AverageMeter() + if mode == "train": + network.train() + elif mode == "valid": + network.eval() else: - logits, logits_aux = logits, None - with torch.no_grad(): - teacher_f, teacher_logits = teacher(inputs) + raise ValueError("The mode is not right : {:}".format(mode)) + teacher.eval() - loss = loss_KD_fn(criterion, logits, teacher_logits, student_f, teacher_f, targets, config.KD_alpha, config.KD_temperature) - if config is not None and hasattr(config, 'auxiliary') and config.auxiliary > 0: - loss_aux = criterion(logits_aux, targets) - loss += config.auxiliary * loss_aux - - if mode == 'train': - loss.backward() - optimizer.step() - - # record - sprec1, sprec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - losses.update(loss.item(), inputs.size(0)) - top1.update (sprec1.item(), inputs.size(0)) - top5.update (sprec5.item(), inputs.size(0)) - # teacher - tprec1, tprec5 = obtain_accuracy(teacher_logits.data, targets.data, topk=(1, 5)) - Ttop1.update (tprec1.item(), inputs.size(0)) - Ttop5.update (tprec5.item(), inputs.size(0)) - - # measure elapsed time - batch_time.update(time.time() - end) + logger.log( + "[{:5s}] config :: auxiliary={:}, KD :: [alpha={:.2f}, temperature={:.2f}]".format( + mode, config.auxiliary if hasattr(config, "auxiliary") else -1, config.KD_alpha, config.KD_temperature + ) + ) end = time.time() + for i, (inputs, targets) in enumerate(xloader): + if mode == "train": + scheduler.update(None, 1.0 * i / len(xloader)) + # measure data loading time + data_time.update(time.time() - end) + # calculate prediction and loss + targets = targets.cuda(non_blocking=True) - if i % print_freq == 0 or (i+1) == len(xloader): - Sstr = ' {:5s} '.format(mode.upper()) + time_string() + ' [{:}][{:03d}/{:03d}]'.format(extra_info, i, len(xloader)) - if scheduler is not None: - Sstr += ' {:}'.format(scheduler.get_min_info()) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Lstr = 'Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})'.format(loss=losses, top1=top1, top5=top5) - Lstr+= ' Teacher : acc@1={:.2f}, acc@5={:.2f}'.format(Ttop1.avg, Ttop5.avg) - Istr = 'Size={:}'.format(list(inputs.size())) - logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Istr) + if mode == "train": + optimizer.zero_grad() - logger.log(' **{:5s}** accuracy drop :: @1={:.2f}, @5={:.2f}'.format(mode.upper(), Ttop1.avg - top1.avg, Ttop5.avg - top5.avg)) - logger.log(' **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}'.format(mode=mode.upper(), top1=top1, top5=top5, error1=100-top1.avg, error5=100-top5.avg, loss=losses.avg)) - return losses.avg, top1.avg, top5.avg + student_f, logits = network(inputs) + if isinstance(logits, list): + assert len(logits) == 2, "logits must has {:} items instead of {:}".format(2, len(logits)) + logits, logits_aux = logits + else: + logits, logits_aux = logits, None + with torch.no_grad(): + teacher_f, teacher_logits = teacher(inputs) + + loss = loss_KD_fn( + criterion, logits, teacher_logits, student_f, teacher_f, targets, config.KD_alpha, config.KD_temperature + ) + if config is not None and hasattr(config, "auxiliary") and config.auxiliary > 0: + loss_aux = criterion(logits_aux, targets) + loss += config.auxiliary * loss_aux + + if mode == "train": + loss.backward() + optimizer.step() + + # record + sprec1, sprec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) + losses.update(loss.item(), inputs.size(0)) + top1.update(sprec1.item(), inputs.size(0)) + top5.update(sprec5.item(), inputs.size(0)) + # teacher + tprec1, tprec5 = obtain_accuracy(teacher_logits.data, targets.data, topk=(1, 5)) + Ttop1.update(tprec1.item(), inputs.size(0)) + Ttop5.update(tprec5.item(), inputs.size(0)) + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + + if i % print_freq == 0 or (i + 1) == len(xloader): + Sstr = ( + " {:5s} ".format(mode.upper()) + + time_string() + + " [{:}][{:03d}/{:03d}]".format(extra_info, i, len(xloader)) + ) + if scheduler is not None: + Sstr += " {:}".format(scheduler.get_min_info()) + Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format( + batch_time=batch_time, data_time=data_time + ) + Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format( + loss=losses, top1=top1, top5=top5 + ) + Lstr += " Teacher : acc@1={:.2f}, acc@5={:.2f}".format(Ttop1.avg, Ttop5.avg) + Istr = "Size={:}".format(list(inputs.size())) + logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr) + + logger.log( + " **{:5s}** accuracy drop :: @1={:.2f}, @5={:.2f}".format( + mode.upper(), Ttop1.avg - top1.avg, Ttop5.avg - top5.avg + ) + ) + logger.log( + " **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}".format( + mode=mode.upper(), top1=top1, top5=top5, error1=100 - top1.avg, error5=100 - top5.avg, loss=losses.avg + ) + ) + return losses.avg, top1.avg, top5.avg diff --git a/lib/procedures/starts.py b/lib/procedures/starts.py index b1b19d3..c93b968 100644 --- a/lib/procedures/starts.py +++ b/lib/procedures/starts.py @@ -3,62 +3,71 @@ ################################################## import os, sys, torch, random, PIL, copy, numpy as np from os import path as osp -from shutil import copyfile +from shutil import copyfile def prepare_seed(rand_seed): - random.seed(rand_seed) - np.random.seed(rand_seed) - torch.manual_seed(rand_seed) - torch.cuda.manual_seed(rand_seed) - torch.cuda.manual_seed_all(rand_seed) + random.seed(rand_seed) + np.random.seed(rand_seed) + torch.manual_seed(rand_seed) + torch.cuda.manual_seed(rand_seed) + torch.cuda.manual_seed_all(rand_seed) def prepare_logger(xargs): - args = copy.deepcopy( xargs ) - from log_utils import Logger - logger = Logger(args.save_dir, args.rand_seed) - logger.log('Main Function with logger : {:}'.format(logger)) - logger.log('Arguments : -------------------------------') - for name, value in args._get_kwargs(): - logger.log('{:16} : {:}'.format(name, value)) - logger.log("Python Version : {:}".format(sys.version.replace('\n', ' '))) - logger.log("Pillow Version : {:}".format(PIL.__version__)) - logger.log("PyTorch Version : {:}".format(torch.__version__)) - logger.log("cuDNN Version : {:}".format(torch.backends.cudnn.version())) - logger.log("CUDA available : {:}".format(torch.cuda.is_available())) - logger.log("CUDA GPU numbers : {:}".format(torch.cuda.device_count())) - logger.log("CUDA_VISIBLE_DEVICES : {:}".format(os.environ['CUDA_VISIBLE_DEVICES'] if 'CUDA_VISIBLE_DEVICES' in os.environ else 'None')) - return logger + args = copy.deepcopy(xargs) + from log_utils import Logger + + logger = Logger(args.save_dir, args.rand_seed) + logger.log("Main Function with logger : {:}".format(logger)) + logger.log("Arguments : -------------------------------") + for name, value in args._get_kwargs(): + logger.log("{:16} : {:}".format(name, value)) + logger.log("Python Version : {:}".format(sys.version.replace("\n", " "))) + logger.log("Pillow Version : {:}".format(PIL.__version__)) + logger.log("PyTorch Version : {:}".format(torch.__version__)) + logger.log("cuDNN Version : {:}".format(torch.backends.cudnn.version())) + logger.log("CUDA available : {:}".format(torch.cuda.is_available())) + logger.log("CUDA GPU numbers : {:}".format(torch.cuda.device_count())) + logger.log( + "CUDA_VISIBLE_DEVICES : {:}".format( + os.environ["CUDA_VISIBLE_DEVICES"] if "CUDA_VISIBLE_DEVICES" in os.environ else "None" + ) + ) + return logger def get_machine_info(): - info = "Python Version : {:}".format(sys.version.replace('\n', ' ')) - info+= "\nPillow Version : {:}".format(PIL.__version__) - info+= "\nPyTorch Version : {:}".format(torch.__version__) - info+= "\ncuDNN Version : {:}".format(torch.backends.cudnn.version()) - info+= "\nCUDA available : {:}".format(torch.cuda.is_available()) - info+= "\nCUDA GPU numbers : {:}".format(torch.cuda.device_count()) - if 'CUDA_VISIBLE_DEVICES' in os.environ: - info+= "\nCUDA_VISIBLE_DEVICES={:}".format(os.environ['CUDA_VISIBLE_DEVICES']) - else: - info+= "\nDoes not set CUDA_VISIBLE_DEVICES" - return info + info = "Python Version : {:}".format(sys.version.replace("\n", " ")) + info += "\nPillow Version : {:}".format(PIL.__version__) + info += "\nPyTorch Version : {:}".format(torch.__version__) + info += "\ncuDNN Version : {:}".format(torch.backends.cudnn.version()) + info += "\nCUDA available : {:}".format(torch.cuda.is_available()) + info += "\nCUDA GPU numbers : {:}".format(torch.cuda.device_count()) + if "CUDA_VISIBLE_DEVICES" in os.environ: + info += "\nCUDA_VISIBLE_DEVICES={:}".format(os.environ["CUDA_VISIBLE_DEVICES"]) + else: + info += "\nDoes not set CUDA_VISIBLE_DEVICES" + return info def save_checkpoint(state, filename, logger): - if osp.isfile(filename): - if hasattr(logger, 'log'): logger.log('Find {:} exist, delete is at first before saving'.format(filename)) - os.remove(filename) - torch.save(state, filename) - assert osp.isfile(filename), 'save filename : {:} failed, which is not found.'.format(filename) - if hasattr(logger, 'log'): logger.log('save checkpoint into {:}'.format(filename)) - return filename + if osp.isfile(filename): + if hasattr(logger, "log"): + logger.log("Find {:} exist, delete is at first before saving".format(filename)) + os.remove(filename) + torch.save(state, filename) + assert osp.isfile(filename), "save filename : {:} failed, which is not found.".format(filename) + if hasattr(logger, "log"): + logger.log("save checkpoint into {:}".format(filename)) + return filename def copy_checkpoint(src, dst, logger): - if osp.isfile(dst): - if hasattr(logger, 'log'): logger.log('Find {:} exist, delete is at first before saving'.format(dst)) - os.remove(dst) - copyfile(src, dst) - if hasattr(logger, 'log'): logger.log('copy the file from {:} into {:}'.format(src, dst)) + if osp.isfile(dst): + if hasattr(logger, "log"): + logger.log("Find {:} exist, delete is at first before saving".format(dst)) + os.remove(dst) + copyfile(src, dst) + if hasattr(logger, "log"): + logger.log("copy the file from {:} into {:}".format(src, dst)) diff --git a/lib/utils/__init__.py b/lib/utils/__init__.py index 2c36e8a..40fa534 100644 --- a/lib/utils/__init__.py +++ b/lib/utils/__init__.py @@ -1,7 +1,7 @@ from .evaluation_utils import obtain_accuracy -from .gpu_manager import GPUManager -from .flop_benchmark import get_model_infos, count_parameters, count_parameters_in_MB -from .affine_utils import normalize_points, denormalize_points -from .affine_utils import identity2affine, solve2theta, affine2image -from .hash_utils import get_md5_file -from .str_utils import split_str2indexes +from .gpu_manager import GPUManager +from .flop_benchmark import get_model_infos, count_parameters, count_parameters_in_MB +from .affine_utils import normalize_points, denormalize_points +from .affine_utils import identity2affine, solve2theta, affine2image +from .hash_utils import get_md5_file +from .str_utils import split_str2indexes diff --git a/lib/utils/affine_utils.py b/lib/utils/affine_utils.py index b122cb4..c452bcd 100644 --- a/lib/utils/affine_utils.py +++ b/lib/utils/affine_utils.py @@ -1,125 +1,149 @@ # functions for affine transformation -import math, torch +import math +import torch import numpy as np import torch.nn.functional as F + def identity2affine(full=False): - if not full: - parameters = torch.zeros((2,3)) - parameters[0, 0] = parameters[1, 1] = 1 - else: - parameters = torch.zeros((3,3)) - parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1 - return parameters + if not full: + parameters = torch.zeros((2, 3)) + parameters[0, 0] = parameters[1, 1] = 1 + else: + parameters = torch.zeros((3, 3)) + parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1 + return parameters + def normalize_L(x, L): - return -1. + 2. * x / (L-1) + return -1.0 + 2.0 * x / (L - 1) + def denormalize_L(x, L): - return (x + 1.0) / 2.0 * (L-1) + return (x + 1.0) / 2.0 * (L - 1) + def crop2affine(crop_box, W, H): - assert len(crop_box) == 4, 'Invalid crop-box : {:}'.format(crop_box) - parameters = torch.zeros(3,3) - x1, y1 = normalize_L(crop_box[0], W), normalize_L(crop_box[1], H) - x2, y2 = normalize_L(crop_box[2], W), normalize_L(crop_box[3], H) - parameters[0,0] = (x2-x1)/2 - parameters[0,2] = (x2+x1)/2 + assert len(crop_box) == 4, "Invalid crop-box : {:}".format(crop_box) + parameters = torch.zeros(3, 3) + x1, y1 = normalize_L(crop_box[0], W), normalize_L(crop_box[1], H) + x2, y2 = normalize_L(crop_box[2], W), normalize_L(crop_box[3], H) + parameters[0, 0] = (x2 - x1) / 2 + parameters[0, 2] = (x2 + x1) / 2 + + parameters[1, 1] = (y2 - y1) / 2 + parameters[1, 2] = (y2 + y1) / 2 + parameters[2, 2] = 1 + return parameters - parameters[1,1] = (y2-y1)/2 - parameters[1,2] = (y2+y1)/2 - parameters[2,2] = 1 - return parameters def scale2affine(scalex, scaley): - parameters = torch.zeros(3,3) - parameters[0,0] = scalex - parameters[1,1] = scaley - parameters[2,2] = 1 - return parameters - + parameters = torch.zeros(3, 3) + parameters[0, 0] = scalex + parameters[1, 1] = scaley + parameters[2, 2] = 1 + return parameters + + def offset2affine(offx, offy): - parameters = torch.zeros(3,3) - parameters[0,0] = parameters[1,1] = parameters[2,2] = 1 - parameters[0,2] = offx - parameters[1,2] = offy - return parameters + parameters = torch.zeros(3, 3) + parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1 + parameters[0, 2] = offx + parameters[1, 2] = offy + return parameters + def horizontalmirror2affine(): - parameters = torch.zeros(3,3) - parameters[0,0] = -1 - parameters[1,1] = parameters[2,2] = 1 - return parameters + parameters = torch.zeros(3, 3) + parameters[0, 0] = -1 + parameters[1, 1] = parameters[2, 2] = 1 + return parameters + # clockwise rotate image = counterclockwise rotate the rectangle # degree is between [0, 360] def rotate2affine(degree): - assert degree >= 0 and degree <= 360, 'Invalid degree : {:}'.format(degree) - degree = degree / 180 * math.pi - parameters = torch.zeros(3,3) - parameters[0,0] = math.cos(-degree) - parameters[0,1] = -math.sin(-degree) - parameters[1,0] = math.sin(-degree) - parameters[1,1] = math.cos(-degree) - parameters[2,2] = 1 - return parameters + assert degree >= 0 and degree <= 360, "Invalid degree : {:}".format(degree) + degree = degree / 180 * math.pi + parameters = torch.zeros(3, 3) + parameters[0, 0] = math.cos(-degree) + parameters[0, 1] = -math.sin(-degree) + parameters[1, 0] = math.sin(-degree) + parameters[1, 1] = math.cos(-degree) + parameters[2, 2] = 1 + return parameters + # shape is a tuple [H, W] def normalize_points(shape, points): - assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape) - assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), 'points are wrong : {:}'.format(points.shape) - (H, W), points = shape, points.clone() - points[0, :] = normalize_L(points[0,:], W) - points[1, :] = normalize_L(points[1,:], H) - return points + assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format( + shape + ) + assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), "points are wrong : {:}".format(points.shape) + (H, W), points = shape, points.clone() + points[0, :] = normalize_L(points[0, :], W) + points[1, :] = normalize_L(points[1, :], H) + return points + # shape is a tuple [H, W] def normalize_points_batch(shape, points): - assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape) - assert isinstance(points, torch.Tensor) and (points.size(-1) == 2), 'points are wrong : {:}'.format(points.shape) - (H, W), points = shape, points.clone() - x = normalize_L(points[...,0], W) - y = normalize_L(points[...,1], H) - return torch.stack((x,y), dim=-1) + assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format( + shape + ) + assert isinstance(points, torch.Tensor) and (points.size(-1) == 2), "points are wrong : {:}".format(points.shape) + (H, W), points = shape, points.clone() + x = normalize_L(points[..., 0], W) + y = normalize_L(points[..., 1], H) + return torch.stack((x, y), dim=-1) + # shape is a tuple [H, W] def denormalize_points(shape, points): - assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape) - assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), 'points are wrong : {:}'.format(points.shape) - (H, W), points = shape, points.clone() - points[0, :] = denormalize_L(points[0,:], W) - points[1, :] = denormalize_L(points[1,:], H) - return points + assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format( + shape + ) + assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), "points are wrong : {:}".format(points.shape) + (H, W), points = shape, points.clone() + points[0, :] = denormalize_L(points[0, :], W) + points[1, :] = denormalize_L(points[1, :], H) + return points + # shape is a tuple [H, W] def denormalize_points_batch(shape, points): - assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape) - assert isinstance(points, torch.Tensor) and (points.shape[-1] == 2), 'points are wrong : {:}'.format(points.shape) - (H, W), points = shape, points.clone() - x = denormalize_L(points[...,0], W) - y = denormalize_L(points[...,1], H) - return torch.stack((x,y), dim=-1) + assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format( + shape + ) + assert isinstance(points, torch.Tensor) and (points.shape[-1] == 2), "points are wrong : {:}".format(points.shape) + (H, W), points = shape, points.clone() + x = denormalize_L(points[..., 0], W) + y = denormalize_L(points[..., 1], H) + return torch.stack((x, y), dim=-1) + # make target * theta = source def solve2theta(source, target): - source, target = source.clone(), target.clone() - oks = source[2, :] == 1 - assert torch.sum(oks).item() >= 3, 'valid points : {:} is short'.format(oks) - if target.size(0) == 2: target = torch.cat((target, oks.unsqueeze(0).float()), dim=0) - source, target = source[:, oks], target[:, oks] - source, target = source.transpose(1,0), target.transpose(1,0) - assert source.size(1) == target.size(1) == 3 - #X, residual, rank, s = np.linalg.lstsq(target.numpy(), source.numpy()) - #theta = torch.Tensor(X.T[:2, :]) - X_, qr = torch.gels(source, target) - theta = X_[:3, :2].transpose(1, 0) - return theta + source, target = source.clone(), target.clone() + oks = source[2, :] == 1 + assert torch.sum(oks).item() >= 3, "valid points : {:} is short".format(oks) + if target.size(0) == 2: + target = torch.cat((target, oks.unsqueeze(0).float()), dim=0) + source, target = source[:, oks], target[:, oks] + source, target = source.transpose(1, 0), target.transpose(1, 0) + assert source.size(1) == target.size(1) == 3 + # X, residual, rank, s = np.linalg.lstsq(target.numpy(), source.numpy()) + # theta = torch.Tensor(X.T[:2, :]) + X_, qr = torch.gels(source, target) + theta = X_[:3, :2].transpose(1, 0) + return theta + # shape = [H,W] def affine2image(image, theta, shape): - C, H, W = image.size() - theta = theta[:2, :].unsqueeze(0) - grid_size = torch.Size([1, C, shape[0], shape[1]]) - grid = F.affine_grid(theta, grid_size) - affI = F.grid_sample(image.unsqueeze(0), grid, mode='bilinear', padding_mode='border') - return affI.squeeze(0) + C, H, W = image.size() + theta = theta[:2, :].unsqueeze(0) + grid_size = torch.Size([1, C, shape[0], shape[1]]) + grid = F.affine_grid(theta, grid_size) + affI = F.grid_sample(image.unsqueeze(0), grid, mode="bilinear", padding_mode="border") + return affI.squeeze(0) diff --git a/lib/utils/evaluation_utils.py b/lib/utils/evaluation_utils.py index cf853c8..088f318 100644 --- a/lib/utils/evaluation_utils.py +++ b/lib/utils/evaluation_utils.py @@ -1,16 +1,17 @@ import torch + def obtain_accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - maxk = max(topk) - batch_size = target.size(0) + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) - res.append(correct_k.mul_(100.0 / batch_size)) - return res + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) + res.append(correct_k.mul_(100.0 / batch_size)) + return res diff --git a/lib/utils/flop_benchmark.py b/lib/utils/flop_benchmark.py index 41dbf72..ade64d7 100644 --- a/lib/utils/flop_benchmark.py +++ b/lib/utils/flop_benchmark.py @@ -4,191 +4,199 @@ import numpy as np def count_parameters_in_MB(model): - return count_parameters(model, "mb") + return count_parameters(model, "mb") def count_parameters(model_or_parameters, unit="mb"): - if isinstance(model_or_parameters, nn.Module): - counts = np.sum(np.prod(v.size()) for v in model_or_parameters.parameters()) - else: - counts = np.sum(np.prod(v.size()) for v in model_or_parameters) - if unit.lower() == "mb": - counts /= 1e6 - elif unit.lower() == "kb": - counts /= 1e3 - elif unit.lower() == "gb": - counts /= 1e9 - elif unit is not None: - raise ValueError("Unknow unit: {:}".format(unit)) - return counts + if isinstance(model_or_parameters, nn.Module): + counts = np.sum(np.prod(v.size()) for v in model_or_parameters.parameters()) + else: + counts = np.sum(np.prod(v.size()) for v in model_or_parameters) + if unit.lower() == "mb": + counts /= 1e6 + elif unit.lower() == "kb": + counts /= 1e3 + elif unit.lower() == "gb": + counts /= 1e9 + elif unit is not None: + raise ValueError("Unknow unit: {:}".format(unit)) + return counts def get_model_infos(model, shape): - #model = copy.deepcopy( model ) + # model = copy.deepcopy( model ) - model = add_flops_counting_methods(model) - #model = model.cuda() - model.eval() + model = add_flops_counting_methods(model) + # model = model.cuda() + model.eval() - #cache_inputs = torch.zeros(*shape).cuda() - #cache_inputs = torch.zeros(*shape) - cache_inputs = torch.rand(*shape) - if next(model.parameters()).is_cuda: cache_inputs = cache_inputs.cuda() - #print_log('In the calculating function : cache input size : {:}'.format(cache_inputs.size()), log) - with torch.no_grad(): - _____ = model(cache_inputs) - FLOPs = compute_average_flops_cost( model ) / 1e6 - Param = count_parameters_in_MB(model) + # cache_inputs = torch.zeros(*shape).cuda() + # cache_inputs = torch.zeros(*shape) + cache_inputs = torch.rand(*shape) + if next(model.parameters()).is_cuda: + cache_inputs = cache_inputs.cuda() + # print_log('In the calculating function : cache input size : {:}'.format(cache_inputs.size()), log) + with torch.no_grad(): + _____ = model(cache_inputs) + FLOPs = compute_average_flops_cost(model) / 1e6 + Param = count_parameters_in_MB(model) - if hasattr(model, 'auxiliary_param'): - aux_params = count_parameters_in_MB(model.auxiliary_param()) - print ('The auxiliary params of this model is : {:}'.format(aux_params)) - print ('We remove the auxiliary params from the total params ({:}) when counting'.format(Param)) - Param = Param - aux_params - - #print_log('FLOPs : {:} MB'.format(FLOPs), log) - torch.cuda.empty_cache() - model.apply( remove_hook_function ) - return FLOPs, Param + if hasattr(model, "auxiliary_param"): + aux_params = count_parameters_in_MB(model.auxiliary_param()) + print("The auxiliary params of this model is : {:}".format(aux_params)) + print("We remove the auxiliary params from the total params ({:}) when counting".format(Param)) + Param = Param - aux_params + + # print_log('FLOPs : {:} MB'.format(FLOPs), log) + torch.cuda.empty_cache() + model.apply(remove_hook_function) + return FLOPs, Param # ---- Public functions -def add_flops_counting_methods( model ): - model.__batch_counter__ = 0 - add_batch_counter_hook_function( model ) - model.apply( add_flops_counter_variable_or_reset ) - model.apply( add_flops_counter_hook_function ) - return model - +def add_flops_counting_methods(model): + model.__batch_counter__ = 0 + add_batch_counter_hook_function(model) + model.apply(add_flops_counter_variable_or_reset) + model.apply(add_flops_counter_hook_function) + return model def compute_average_flops_cost(model): - """ - A method that will be available after add_flops_counting_methods() is called on a desired net object. - Returns current mean flops consumption per image. - """ - batches_count = model.__batch_counter__ - flops_sum = 0 - #or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \ - for module in model.modules(): - if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear) \ - or isinstance(module, torch.nn.Conv1d) \ - or hasattr(module, 'calculate_flop_self'): - flops_sum += module.__flops__ - return flops_sum / batches_count + """ + A method that will be available after add_flops_counting_methods() is called on a desired net object. + Returns current mean flops consumption per image. + """ + batches_count = model.__batch_counter__ + flops_sum = 0 + # or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \ + for module in model.modules(): + if ( + isinstance(module, torch.nn.Conv2d) + or isinstance(module, torch.nn.Linear) + or isinstance(module, torch.nn.Conv1d) + or hasattr(module, "calculate_flop_self") + ): + flops_sum += module.__flops__ + return flops_sum / batches_count # ---- Internal functions def pool_flops_counter_hook(pool_module, inputs, output): - batch_size = inputs[0].size(0) - kernel_size = pool_module.kernel_size - out_C, output_height, output_width = output.shape[1:] - assert out_C == inputs[0].size(1), '{:} vs. {:}'.format(out_C, inputs[0].size()) + batch_size = inputs[0].size(0) + kernel_size = pool_module.kernel_size + out_C, output_height, output_width = output.shape[1:] + assert out_C == inputs[0].size(1), "{:} vs. {:}".format(out_C, inputs[0].size()) - overall_flops = batch_size * out_C * output_height * output_width * kernel_size * kernel_size - pool_module.__flops__ += overall_flops + overall_flops = batch_size * out_C * output_height * output_width * kernel_size * kernel_size + pool_module.__flops__ += overall_flops def self_calculate_flops_counter_hook(self_module, inputs, output): - overall_flops = self_module.calculate_flop_self(inputs[0].shape, output.shape) - self_module.__flops__ += overall_flops + overall_flops = self_module.calculate_flop_self(inputs[0].shape, output.shape) + self_module.__flops__ += overall_flops def fc_flops_counter_hook(fc_module, inputs, output): - batch_size = inputs[0].size(0) - xin, xout = fc_module.in_features, fc_module.out_features - assert xin == inputs[0].size(1) and xout == output.size(1), 'IO=({:}, {:})'.format(xin, xout) - overall_flops = batch_size * xin * xout - if fc_module.bias is not None: - overall_flops += batch_size * xout - fc_module.__flops__ += overall_flops + batch_size = inputs[0].size(0) + xin, xout = fc_module.in_features, fc_module.out_features + assert xin == inputs[0].size(1) and xout == output.size(1), "IO=({:}, {:})".format(xin, xout) + overall_flops = batch_size * xin * xout + if fc_module.bias is not None: + overall_flops += batch_size * xout + fc_module.__flops__ += overall_flops def conv1d_flops_counter_hook(conv_module, inputs, outputs): - batch_size = inputs[0].size(0) - outL = outputs.shape[-1] - [kernel] = conv_module.kernel_size - in_channels = conv_module.in_channels - out_channels = conv_module.out_channels - groups = conv_module.groups - conv_per_position_flops = kernel * in_channels * out_channels / groups - - active_elements_count = batch_size * outL - overall_flops = conv_per_position_flops * active_elements_count + batch_size = inputs[0].size(0) + outL = outputs.shape[-1] + [kernel] = conv_module.kernel_size + in_channels = conv_module.in_channels + out_channels = conv_module.out_channels + groups = conv_module.groups + conv_per_position_flops = kernel * in_channels * out_channels / groups - if conv_module.bias is not None: - overall_flops += out_channels * active_elements_count - conv_module.__flops__ += overall_flops + active_elements_count = batch_size * outL + overall_flops = conv_per_position_flops * active_elements_count + + if conv_module.bias is not None: + overall_flops += out_channels * active_elements_count + conv_module.__flops__ += overall_flops def conv2d_flops_counter_hook(conv_module, inputs, output): - batch_size = inputs[0].size(0) - output_height, output_width = output.shape[2:] - - kernel_height, kernel_width = conv_module.kernel_size - in_channels = conv_module.in_channels - out_channels = conv_module.out_channels - groups = conv_module.groups - conv_per_position_flops = kernel_height * kernel_width * in_channels * out_channels / groups - - active_elements_count = batch_size * output_height * output_width - overall_flops = conv_per_position_flops * active_elements_count - - if conv_module.bias is not None: - overall_flops += out_channels * active_elements_count - conv_module.__flops__ += overall_flops + batch_size = inputs[0].size(0) + output_height, output_width = output.shape[2:] + + kernel_height, kernel_width = conv_module.kernel_size + in_channels = conv_module.in_channels + out_channels = conv_module.out_channels + groups = conv_module.groups + conv_per_position_flops = kernel_height * kernel_width * in_channels * out_channels / groups + + active_elements_count = batch_size * output_height * output_width + overall_flops = conv_per_position_flops * active_elements_count + + if conv_module.bias is not None: + overall_flops += out_channels * active_elements_count + conv_module.__flops__ += overall_flops + - def batch_counter_hook(module, inputs, output): - # Can have multiple inputs, getting the first one - inputs = inputs[0] - batch_size = inputs.shape[0] - module.__batch_counter__ += batch_size + # Can have multiple inputs, getting the first one + inputs = inputs[0] + batch_size = inputs.shape[0] + module.__batch_counter__ += batch_size def add_batch_counter_hook_function(module): - if not hasattr(module, '__batch_counter_handle__'): - handle = module.register_forward_hook(batch_counter_hook) - module.__batch_counter_handle__ = handle + if not hasattr(module, "__batch_counter_handle__"): + handle = module.register_forward_hook(batch_counter_hook) + module.__batch_counter_handle__ = handle + - def add_flops_counter_variable_or_reset(module): - if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear) \ - or isinstance(module, torch.nn.Conv1d) \ - or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \ - or hasattr(module, 'calculate_flop_self'): - module.__flops__ = 0 + if ( + isinstance(module, torch.nn.Conv2d) + or isinstance(module, torch.nn.Linear) + or isinstance(module, torch.nn.Conv1d) + or isinstance(module, torch.nn.AvgPool2d) + or isinstance(module, torch.nn.MaxPool2d) + or hasattr(module, "calculate_flop_self") + ): + module.__flops__ = 0 def add_flops_counter_hook_function(module): - if isinstance(module, torch.nn.Conv2d): - if not hasattr(module, '__flops_handle__'): - handle = module.register_forward_hook(conv2d_flops_counter_hook) - module.__flops_handle__ = handle - elif isinstance(module, torch.nn.Conv1d): - if not hasattr(module, '__flops_handle__'): - handle = module.register_forward_hook(conv1d_flops_counter_hook) - module.__flops_handle__ = handle - elif isinstance(module, torch.nn.Linear): - if not hasattr(module, '__flops_handle__'): - handle = module.register_forward_hook(fc_flops_counter_hook) - module.__flops_handle__ = handle - elif isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d): - if not hasattr(module, '__flops_handle__'): - handle = module.register_forward_hook(pool_flops_counter_hook) - module.__flops_handle__ = handle - elif hasattr(module, 'calculate_flop_self'): # self-defined module - if not hasattr(module, '__flops_handle__'): - handle = module.register_forward_hook(self_calculate_flops_counter_hook) - module.__flops_handle__ = handle + if isinstance(module, torch.nn.Conv2d): + if not hasattr(module, "__flops_handle__"): + handle = module.register_forward_hook(conv2d_flops_counter_hook) + module.__flops_handle__ = handle + elif isinstance(module, torch.nn.Conv1d): + if not hasattr(module, "__flops_handle__"): + handle = module.register_forward_hook(conv1d_flops_counter_hook) + module.__flops_handle__ = handle + elif isinstance(module, torch.nn.Linear): + if not hasattr(module, "__flops_handle__"): + handle = module.register_forward_hook(fc_flops_counter_hook) + module.__flops_handle__ = handle + elif isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d): + if not hasattr(module, "__flops_handle__"): + handle = module.register_forward_hook(pool_flops_counter_hook) + module.__flops_handle__ = handle + elif hasattr(module, "calculate_flop_self"): # self-defined module + if not hasattr(module, "__flops_handle__"): + handle = module.register_forward_hook(self_calculate_flops_counter_hook) + module.__flops_handle__ = handle def remove_hook_function(module): - hookers = ['__batch_counter_handle__', '__flops_handle__'] - for hooker in hookers: - if hasattr(module, hooker): - handle = getattr(module, hooker) - handle.remove() - keys = ['__flops__', '__batch_counter__', '__flops__'] + hookers - for ckey in keys: - if hasattr(module, ckey): delattr(module, ckey) + hookers = ["__batch_counter_handle__", "__flops_handle__"] + for hooker in hookers: + if hasattr(module, hooker): + handle = getattr(module, hooker) + handle.remove() + keys = ["__flops__", "__batch_counter__", "__flops__"] + hookers + for ckey in keys: + if hasattr(module, ckey): + delattr(module, ckey) diff --git a/lib/utils/gpu_manager.py b/lib/utils/gpu_manager.py index 520cff4..044c64c 100644 --- a/lib/utils/gpu_manager.py +++ b/lib/utils/gpu_manager.py @@ -1,65 +1,69 @@ import os -class GPUManager(): - queries = ('index', 'gpu_name', 'memory.free', 'memory.used', 'memory.total', 'power.draw', 'power.limit') - def __init__(self): - all_gpus = self.query_gpu(False) +class GPUManager: + queries = ("index", "gpu_name", "memory.free", "memory.used", "memory.total", "power.draw", "power.limit") - def get_info(self, ctype): - cmd = 'nvidia-smi --query-gpu={} --format=csv,noheader'.format(ctype) - lines = os.popen(cmd).readlines() - lines = [line.strip('\n') for line in lines] - return lines + def __init__(self): + all_gpus = self.query_gpu(False) - def query_gpu(self, show=True): - num_gpus = len( self.get_info('index') ) - all_gpus = [ {} for i in range(num_gpus) ] - for query in self.queries: - infos = self.get_info(query) - for idx, info in enumerate(infos): - all_gpus[idx][query] = info + def get_info(self, ctype): + cmd = "nvidia-smi --query-gpu={} --format=csv,noheader".format(ctype) + lines = os.popen(cmd).readlines() + lines = [line.strip("\n") for line in lines] + return lines - if 'CUDA_VISIBLE_DEVICES' in os.environ: - CUDA_VISIBLE_DEVICES = os.environ['CUDA_VISIBLE_DEVICES'].split(',') - selected_gpus = [] - for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES): - find = False - for gpu in all_gpus: - if gpu['index'] == CUDA_VISIBLE_DEVICE: - assert not find, 'Duplicate cuda device index : {}'.format(CUDA_VISIBLE_DEVICE) - find = True - selected_gpus.append( gpu.copy() ) - selected_gpus[-1]['index'] = '{}'.format(idx) - assert find, 'Does not find the device : {}'.format(CUDA_VISIBLE_DEVICE) - all_gpus = selected_gpus - - if show: - allstrings = '' - for gpu in all_gpus: - string = '| ' + def query_gpu(self, show=True): + num_gpus = len(self.get_info("index")) + all_gpus = [{} for i in range(num_gpus)] for query in self.queries: - if query.find('memory') == 0: xinfo = '{:>9}'.format(gpu[query]) - else: xinfo = gpu[query] - string = string + query + ' : ' + xinfo + ' | ' - allstrings = allstrings + string + '\n' - return allstrings - else: - return all_gpus + infos = self.get_info(query) + for idx, info in enumerate(infos): + all_gpus[idx][query] = info + + if "CUDA_VISIBLE_DEVICES" in os.environ: + CUDA_VISIBLE_DEVICES = os.environ["CUDA_VISIBLE_DEVICES"].split(",") + selected_gpus = [] + for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES): + find = False + for gpu in all_gpus: + if gpu["index"] == CUDA_VISIBLE_DEVICE: + assert not find, "Duplicate cuda device index : {}".format(CUDA_VISIBLE_DEVICE) + find = True + selected_gpus.append(gpu.copy()) + selected_gpus[-1]["index"] = "{}".format(idx) + assert find, "Does not find the device : {}".format(CUDA_VISIBLE_DEVICE) + all_gpus = selected_gpus + + if show: + allstrings = "" + for gpu in all_gpus: + string = "| " + for query in self.queries: + if query.find("memory") == 0: + xinfo = "{:>9}".format(gpu[query]) + else: + xinfo = gpu[query] + string = string + query + " : " + xinfo + " | " + allstrings = allstrings + string + "\n" + return allstrings + else: + return all_gpus + + def select_by_memory(self, numbers=1): + all_gpus = self.query_gpu(False) + assert numbers <= len(all_gpus), "Require {} gpus more than you have".format(numbers) + alls = [] + for idx, gpu in enumerate(all_gpus): + free_memory = gpu["memory.free"] + free_memory = free_memory.split(" ")[0] + free_memory = int(free_memory) + index = gpu["index"] + alls.append((free_memory, index)) + alls.sort(reverse=True) + alls = [int(alls[i][1]) for i in range(numbers)] + return sorted(alls) - def select_by_memory(self, numbers=1): - all_gpus = self.query_gpu(False) - assert numbers <= len(all_gpus), 'Require {} gpus more than you have'.format(numbers) - alls = [] - for idx, gpu in enumerate(all_gpus): - free_memory = gpu['memory.free'] - free_memory = free_memory.split(' ')[0] - free_memory = int(free_memory) - index = gpu['index'] - alls.append((free_memory, index)) - alls.sort(reverse = True) - alls = [ int(alls[i][1]) for i in range(numbers) ] - return sorted(alls) """ if __name__ == '__main__': diff --git a/lib/utils/hash_utils.py b/lib/utils/hash_utils.py index b372cac..53696ae 100644 --- a/lib/utils/hash_utils.py +++ b/lib/utils/hash_utils.py @@ -1,16 +1,17 @@ -import os, hashlib +import os +import hashlib def get_md5_file(file_path, post_truncated=5): - md5_hash = hashlib.md5() - if os.path.exists(file_path): - xfile = open(file_path, "rb") - content = xfile.read() - md5_hash.update(content) - digest = md5_hash.hexdigest() - else: - raise ValueError('[get_md5_file] {:} does not exist'.format(file_path)) - if post_truncated is None: - return digest - else: - return digest[-post_truncated:] + md5_hash = hashlib.md5() + if os.path.exists(file_path): + xfile = open(file_path, "rb") + content = xfile.read() + md5_hash.update(content) + digest = md5_hash.hexdigest() + else: + raise ValueError("[get_md5_file] {:} does not exist".format(file_path)) + if post_truncated is None: + return digest + else: + return digest[-post_truncated:] diff --git a/lib/utils/nas_utils.py b/lib/utils/nas_utils.py index 1b1a44d..e21da1f 100644 --- a/lib/utils/nas_utils.py +++ b/lib/utils/nas_utils.py @@ -10,48 +10,58 @@ from log_utils import time_string def evaluate_one_shot(model, xloader, api, cal_mode, seed=111): - print ('This is an old version of codes to use NAS-Bench-API, and should be modified to align with the new version. Please contact me for more details if you use this function.') - weights = deepcopy(model.state_dict()) - model.train(cal_mode) - with torch.no_grad(): - logits = nn.functional.log_softmax(model.arch_parameters, dim=-1) - archs = CellStructure.gen_all(model.op_names, model.max_nodes, False) - probs, accuracies, gt_accs_10_valid, gt_accs_10_test = [], [], [], [] - loader_iter = iter(xloader) - random.seed(seed) - random.shuffle(archs) - for idx, arch in enumerate(archs): - arch_index = api.query_index_by_arch( arch ) - metrics = api.get_more_info(arch_index, 'cifar10-valid', None, False, False) - gt_accs_10_valid.append( metrics['valid-accuracy'] ) - metrics = api.get_more_info(arch_index, 'cifar10', None, False, False) - gt_accs_10_test.append( metrics['test-accuracy'] ) - select_logits = [] - for i, node_info in enumerate(arch.nodes): - for op, xin in node_info: - node_str = '{:}<-{:}'.format(i+1, xin) - op_index = model.op_names.index(op) - select_logits.append( logits[model.edge2index[node_str], op_index] ) - cur_prob = sum(select_logits).item() - probs.append( cur_prob ) - cor_prob_valid = np.corrcoef(probs, gt_accs_10_valid)[0,1] - cor_prob_test = np.corrcoef(probs, gt_accs_10_test )[0,1] - print ('{:} correlation for probabilities : {:.6f} on CIFAR-10 validation and {:.6f} on CIFAR-10 test'.format(time_string(), cor_prob_valid, cor_prob_test)) - - for idx, arch in enumerate(archs): - model.set_cal_mode('dynamic', arch) - try: - inputs, targets = next(loader_iter) - except: + print( + "This is an old version of codes to use NAS-Bench-API, and should be modified to align with the new version. Please contact me for more details if you use this function." + ) + weights = deepcopy(model.state_dict()) + model.train(cal_mode) + with torch.no_grad(): + logits = nn.functional.log_softmax(model.arch_parameters, dim=-1) + archs = CellStructure.gen_all(model.op_names, model.max_nodes, False) + probs, accuracies, gt_accs_10_valid, gt_accs_10_test = [], [], [], [] loader_iter = iter(xloader) - inputs, targets = next(loader_iter) - _, logits = model(inputs.cuda()) - _, preds = torch.max(logits, dim=-1) - correct = (preds == targets.cuda() ).float() - accuracies.append( correct.mean().item() ) - if idx != 0 and (idx % 500 == 0 or idx + 1 == len(archs)): - cor_accs_valid = np.corrcoef(accuracies, gt_accs_10_valid[:idx+1])[0,1] - cor_accs_test = np.corrcoef(accuracies, gt_accs_10_test [:idx+1])[0,1] - print ('{:} {:05d}/{:05d} mode={:5s}, correlation : accs={:.5f} for CIFAR-10 valid, {:.5f} for CIFAR-10 test.'.format(time_string(), idx, len(archs), 'Train' if cal_mode else 'Eval', cor_accs_valid, cor_accs_test)) - model.load_state_dict(weights) - return archs, probs, accuracies + random.seed(seed) + random.shuffle(archs) + for idx, arch in enumerate(archs): + arch_index = api.query_index_by_arch(arch) + metrics = api.get_more_info(arch_index, "cifar10-valid", None, False, False) + gt_accs_10_valid.append(metrics["valid-accuracy"]) + metrics = api.get_more_info(arch_index, "cifar10", None, False, False) + gt_accs_10_test.append(metrics["test-accuracy"]) + select_logits = [] + for i, node_info in enumerate(arch.nodes): + for op, xin in node_info: + node_str = "{:}<-{:}".format(i + 1, xin) + op_index = model.op_names.index(op) + select_logits.append(logits[model.edge2index[node_str], op_index]) + cur_prob = sum(select_logits).item() + probs.append(cur_prob) + cor_prob_valid = np.corrcoef(probs, gt_accs_10_valid)[0, 1] + cor_prob_test = np.corrcoef(probs, gt_accs_10_test)[0, 1] + print( + "{:} correlation for probabilities : {:.6f} on CIFAR-10 validation and {:.6f} on CIFAR-10 test".format( + time_string(), cor_prob_valid, cor_prob_test + ) + ) + + for idx, arch in enumerate(archs): + model.set_cal_mode("dynamic", arch) + try: + inputs, targets = next(loader_iter) + except: + loader_iter = iter(xloader) + inputs, targets = next(loader_iter) + _, logits = model(inputs.cuda()) + _, preds = torch.max(logits, dim=-1) + correct = (preds == targets.cuda()).float() + accuracies.append(correct.mean().item()) + if idx != 0 and (idx % 500 == 0 or idx + 1 == len(archs)): + cor_accs_valid = np.corrcoef(accuracies, gt_accs_10_valid[: idx + 1])[0, 1] + cor_accs_test = np.corrcoef(accuracies, gt_accs_10_test[: idx + 1])[0, 1] + print( + "{:} {:05d}/{:05d} mode={:5s}, correlation : accs={:.5f} for CIFAR-10 valid, {:.5f} for CIFAR-10 test.".format( + time_string(), idx, len(archs), "Train" if cal_mode else "Eval", cor_accs_valid, cor_accs_test + ) + ) + model.load_state_dict(weights) + return archs, probs, accuracies diff --git a/lib/utils/str_utils.py b/lib/utils/str_utils.py index bd58b6a..7e49fe1 100644 --- a/lib/utils/str_utils.py +++ b/lib/utils/str_utils.py @@ -1,18 +1,17 @@ - def split_str2indexes(string: str, max_check: int, length_limit=5): - if not isinstance(string, str): - raise ValueError('Invalid scheme for {:}'.format(string)) - srangestr = "".join(string.split()) - indexes = set() - for srange in srangestr.split(','): - srange = srange.split('-') - if len(srange) != 2: - raise ValueError('invalid srange : {:}'.format(srange)) - if length_limit is not None: - assert len(srange[0]) == len(srange[1]) == length_limit, 'invalid srange : {:}'.format(srange) - srange = (int(srange[0]), int(srange[1])) - if not (0 <= srange[0] <= srange[1] < max_check): - raise ValueError('{:} vs {:} vs {:}'.format(srange[0], srange[1], max_check)) - for i in range(srange[0], srange[1]+1): - indexes.add(i) - return indexes + if not isinstance(string, str): + raise ValueError("Invalid scheme for {:}".format(string)) + srangestr = "".join(string.split()) + indexes = set() + for srange in srangestr.split(","): + srange = srange.split("-") + if len(srange) != 2: + raise ValueError("invalid srange : {:}".format(srange)) + if length_limit is not None: + assert len(srange[0]) == len(srange[1]) == length_limit, "invalid srange : {:}".format(srange) + srange = (int(srange[0]), int(srange[1])) + if not (0 <= srange[0] <= srange[1] < max_check): + raise ValueError("{:} vs {:} vs {:}".format(srange[0], srange[1], max_check)) + for i in range(srange[0], srange[1] + 1): + indexes.add(i) + return indexes diff --git a/lib/utils/weight_watcher.py b/lib/utils/weight_watcher.py index a132c44..08404c4 100644 --- a/lib/utils/weight_watcher.py +++ b/lib/utils/weight_watcher.py @@ -11,309 +11,350 @@ from sklearn.decomposition import TruncatedSVD def available_module_types(): - return (nn.Conv2d, nn.Linear) + return (nn.Conv2d, nn.Linear) def get_conv2D_Wmats(tensor: np.ndarray) -> List[np.ndarray]: - """ - Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j). - Return ij (N x M) matrices - """ - mats = [] - N, M, imax, jmax = tensor.shape - assert N + M >= imax + jmax, 'invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)'.format(N, M, imax, jmax) - for i in range(imax): - for j in range(jmax): - w = tensor[:, :, i, j] - if N < M: w = w.T - mats.append(w) - return mats + """ + Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j). + Return ij (N x M) matrices + """ + mats = [] + N, M, imax, jmax = tensor.shape + assert N + M >= imax + jmax, "invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)".format(N, M, imax, jmax) + for i in range(imax): + for j in range(jmax): + w = tensor[:, :, i, j] + if N < M: + w = w.T + mats.append(w) + return mats def glorot_norm_check(W, N, M, rf_size, lower=0.5, upper=1.5): - """Check if this layer needs Glorot Normalization Fix""" + """Check if this layer needs Glorot Normalization Fix""" - kappa = np.sqrt(2 / ((N + M) * rf_size)) - norm = np.linalg.norm(W) + kappa = np.sqrt(2 / ((N + M) * rf_size)) + norm = np.linalg.norm(W) - check1 = norm / np.sqrt(N * M) - check2 = norm / (kappa * np.sqrt(N * M)) + check1 = norm / np.sqrt(N * M) + check2 = norm / (kappa * np.sqrt(N * M)) + + if (rf_size > 1) and (check2 > lower) and (check2 < upper): + return check2, True + elif (check1 > lower) & (check1 < upper): + return check1, True + else: + if rf_size > 1: + return check2, False + else: + return check1, False - if (rf_size > 1) and (check2 > lower) and (check2 < upper): - return check2, True - elif (check1 > lower) & (check1 < upper): - return check1, True - else: - if rf_size > 1: return check2, False - else: return check1, False def glorot_norm_fix(w, n, m, rf_size): - """Apply Glorot Normalization Fix.""" - kappa = np.sqrt(2 / ((n + m) * rf_size)) - w = w / kappa - return w + """Apply Glorot Normalization Fix.""" + kappa = np.sqrt(2 / ((n + m) * rf_size)) + w = w / kappa + return w def analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix): - results = OrderedDict() - count = len(weights) - if count == 0: return results + results = OrderedDict() + count = len(weights) + if count == 0: + return results - for i, weight in enumerate(weights): - M, N = np.min(weight.shape), np.max(weight.shape) - Q = N / M - results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q) - check, checkTF = glorot_norm_check(weight, N, M, count) - cur_res['check'] = check - cur_res['checkTF'] = checkTF - # assume receptive field size is count - if glorot_fix: - weight = glorot_norm_fix(weight, N, M, count) - else: - # probably never needed since we always fix for glorot - weight = weight * np.sqrt(count / 2.0) + for i, weight in enumerate(weights): + M, N = np.min(weight.shape), np.max(weight.shape) + Q = N / M + results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q) + check, checkTF = glorot_norm_check(weight, N, M, count) + cur_res["check"] = check + cur_res["checkTF"] = checkTF + # assume receptive field size is count + if glorot_fix: + weight = glorot_norm_fix(weight, N, M, count) + else: + # probably never needed since we always fix for glorot + weight = weight * np.sqrt(count / 2.0) - if spectralnorms: # spectralnorm is the max eigenvalues - svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10) - svd.fit(weight) - sv = svd.singular_values_ - sv_max = np.max(sv) - if normalize: - evals = sv * sv / N - else: - evals = sv * sv - lambda0 = evals[0] - cur_res["spectralnorm"] = lambda0 - cur_res["logspectralnorm"] = np.log10(lambda0) - else: - lambda0 = None + if spectralnorms: # spectralnorm is the max eigenvalues + svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10) + svd.fit(weight) + sv = svd.singular_values_ + sv_max = np.max(sv) + if normalize: + evals = sv * sv / N + else: + evals = sv * sv + lambda0 = evals[0] + cur_res["spectralnorm"] = lambda0 + cur_res["logspectralnorm"] = np.log10(lambda0) + else: + lambda0 = None - if M < min_size: - summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(i + 1, count, M, N, min_size) - cur_res["summary"] = summary - continue - elif max_size > 0 and M > max_size: - summary = "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format(i + 1, count, M, N, max_size) - cur_res["summary"] = summary - continue - else: - summary = [] - if alphas: - import powerlaw - svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10) - svd.fit(weight.astype(float)) - sv = svd.singular_values_ - if normalize: evals = sv * sv / N - else: evals = sv * sv + if M < min_size: + summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(i + 1, count, M, N, min_size) + cur_res["summary"] = summary + continue + elif max_size > 0 and M > max_size: + summary = "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format( + i + 1, count, M, N, max_size + ) + cur_res["summary"] = summary + continue + else: + summary = [] + if alphas: + import powerlaw - lambda_max = np.max(evals) - fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False) - alpha = fit.alpha - cur_res["alpha"] = alpha - D = fit.D - cur_res["D"] = D - cur_res["lambda_min"] = np.min(evals) - cur_res["lambda_max"] = lambda_max - alpha_weighted = alpha * np.log10(lambda_max) - cur_res["alpha_weighted"] = alpha_weighted - tolerance = lambda_max * M * np.finfo(np.max(sv)).eps - cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1) + svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10) + svd.fit(weight.astype(float)) + sv = svd.singular_values_ + if normalize: + evals = sv * sv / N + else: + evals = sv * sv - logpnorm = np.log10(np.sum([ev ** alpha for ev in evals])) - cur_res["logpnorm"] = logpnorm + lambda_max = np.max(evals) + fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False) + alpha = fit.alpha + cur_res["alpha"] = alpha + D = fit.D + cur_res["D"] = D + cur_res["lambda_min"] = np.min(evals) + cur_res["lambda_max"] = lambda_max + alpha_weighted = alpha * np.log10(lambda_max) + cur_res["alpha_weighted"] = alpha_weighted + tolerance = lambda_max * M * np.finfo(np.max(sv)).eps + cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1) - summary.append( - "Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format(i + 1, count, M, N, alpha, - alpha_weighted, D, - logpnorm)) + logpnorm = np.log10(np.sum([ev ** alpha for ev in evals])) + cur_res["logpnorm"] = logpnorm - if lognorms: - norm = np.linalg.norm(weight) # Frobenius Norm - cur_res["norm"] = norm - lognorm = np.log10(norm) - cur_res["lognorm"] = lognorm + summary.append( + "Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format( + i + 1, count, M, N, alpha, alpha_weighted, D, logpnorm + ) + ) - X = np.dot(weight.T, weight) - if normalize: X = X / N - normX = np.linalg.norm(X) # Frobenius Norm - cur_res["normX"] = normX - lognormX = np.log10(normX) - cur_res["lognormX"] = lognormX + if lognorms: + norm = np.linalg.norm(weight) # Frobenius Norm + cur_res["norm"] = norm + lognorm = np.log10(norm) + cur_res["lognorm"] = lognorm - summary.append( - "Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(i + 1, count, M, N, lognorm, lognormX)) + X = np.dot(weight.T, weight) + if normalize: + X = X / N + normX = np.linalg.norm(X) # Frobenius Norm + cur_res["normX"] = normX + lognormX = np.log10(normX) + cur_res["lognormX"] = lognormX - if softranks: - softrank = norm ** 2 / sv_max ** 2 - softranklog = np.log10(softrank) - softranklogratio = lognorm / np.log10(sv_max) - cur_res["softrank"] = softrank - cur_res["softranklog"] = softranklog - cur_res["softranklogratio"] = softranklogratio - summary += "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format(summary, softrank, softranklog, - softranklogratio) - cur_res["summary"] = "\n".join(summary) - return results + summary.append( + "Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(i + 1, count, M, N, lognorm, lognormX) + ) + + if softranks: + softrank = norm ** 2 / sv_max ** 2 + softranklog = np.log10(softrank) + softranklogratio = lognorm / np.log10(sv_max) + cur_res["softrank"] = softrank + cur_res["softranklog"] = softranklog + cur_res["softranklogratio"] = softranklogratio + summary += "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format( + summary, softrank, softranklog, softranklogratio + ) + cur_res["summary"] = "\n".join(summary) + return results def compute_details(results): - """ - Return a pandas data frame. - """ - final_summary = OrderedDict() + """ + Return a pandas data frame. + """ + final_summary = OrderedDict() - metrics = { - # key in "results" : pretty print name - "check": "Check", - "checkTF": "CheckTF", - "norm": "Norm", - "lognorm": "LogNorm", - "normX": "Norm X", - "lognormX": "LogNorm X", - "alpha": "Alpha", - "alpha_weighted": "Alpha Weighted", - "spectralnorm": "Spectral Norm", - "logspectralnorm": "Log Spectral Norm", - "softrank": "Softrank", - "softranklog": "Softrank Log", - "softranklogratio": "Softrank Log Ratio", - "sigma_mp": "Marchenko-Pastur (MP) fit sigma", - "numofSpikes": "Number of spikes per MP fit", - "ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals", - "softrank_mp": "Softrank for MP fit", - "logpnorm": "alpha pNorm" - } + metrics = { + # key in "results" : pretty print name + "check": "Check", + "checkTF": "CheckTF", + "norm": "Norm", + "lognorm": "LogNorm", + "normX": "Norm X", + "lognormX": "LogNorm X", + "alpha": "Alpha", + "alpha_weighted": "Alpha Weighted", + "spectralnorm": "Spectral Norm", + "logspectralnorm": "Log Spectral Norm", + "softrank": "Softrank", + "softranklog": "Softrank Log", + "softranklogratio": "Softrank Log Ratio", + "sigma_mp": "Marchenko-Pastur (MP) fit sigma", + "numofSpikes": "Number of spikes per MP fit", + "ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals", + "softrank_mp": "Softrank for MP fit", + "logpnorm": "alpha pNorm", + } - metrics_stats = [] - for metric in metrics: - metrics_stats.append("{}_min".format(metric)) - metrics_stats.append("{}_max".format(metric)) - metrics_stats.append("{}_avg".format(metric)) - - metrics_stats.append("{}_compound_min".format(metric)) - metrics_stats.append("{}_compound_max".format(metric)) - metrics_stats.append("{}_compound_avg".format(metric)) - - columns = ["layer_id", "layer_type", "N", "M", "layer_count", "slice", - "slice_count", "level", "comment"] + [*metrics] + metrics_stats - - metrics_values = {} - metrics_values_compound = {} - - for metric in metrics: - metrics_values[metric] = [] - metrics_values_compound[metric] = [] - - layer_count = 0 - for layer_id, result in results.items(): - layer_count += 1 - - layer_type = np.NAN - if "layer_type" in result: - layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "") - - compounds = {} # temp var + metrics_stats = [] for metric in metrics: - compounds[metric] = [] + metrics_stats.append("{}_min".format(metric)) + metrics_stats.append("{}_max".format(metric)) + metrics_stats.append("{}_avg".format(metric)) - slice_count, Ntotal, Mtotal = 0, 0, 0 - for slice_id, summary in result.items(): - if not str(slice_id).isdigit(): - continue - slice_count += 1 + metrics_stats.append("{}_compound_min".format(metric)) + metrics_stats.append("{}_compound_max".format(metric)) + metrics_stats.append("{}_compound_avg".format(metric)) - N = np.NAN - if "N" in summary: - N = summary["N"] - Ntotal += N + columns = ( + ["layer_id", "layer_type", "N", "M", "layer_count", "slice", "slice_count", "level", "comment"] + + [*metrics] + + metrics_stats + ) - M = np.NAN - if "M" in summary: - M = summary["M"] - Mtotal += M + metrics_values = {} + metrics_values_compound = {} - data = {"layer_id": layer_id, "layer_type": layer_type, "N": N, "M": M, "slice": slice_id, "level": "SLICE", - "comment": "Slice level"} - for metric in metrics: - if metric in summary: - value = summary[metric] - if value is not None: - metrics_values[metric].append(value) - compounds[metric].append(value) - data[metric] = value + for metric in metrics: + metrics_values[metric] = [] + metrics_values_compound[metric] = [] - data = {"layer_id": layer_id, "layer_type": layer_type, "N": Ntotal, "M": Mtotal, "slice_count": slice_count, - "level": "LAYER", "comment": "Layer level"} - # Compute the compound value over the slices - for metric, value in compounds.items(): - count = len(value) - if count == 0: - continue + layer_count = 0 + for layer_id, result in results.items(): + layer_count += 1 - compound = np.mean(value) - metrics_values_compound[metric].append(compound) - data[metric] = compound + layer_type = np.NAN + if "layer_type" in result: + layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "") - data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"} - for metric, metric_name in metrics.items(): - if metric not in metrics_values or len(metrics_values[metric]) == 0: - continue + compounds = {} # temp var + for metric in metrics: + compounds[metric] = [] - values = metrics_values[metric] - minimum = min(values) - maximum = max(values) - avg = np.mean(values) - final_summary[metric] = avg - # print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) - data["{}_min".format(metric)] = minimum - data["{}_max".format(metric)] = maximum - data["{}_avg".format(metric)] = avg + slice_count, Ntotal, Mtotal = 0, 0, 0 + for slice_id, summary in result.items(): + if not str(slice_id).isdigit(): + continue + slice_count += 1 - values = metrics_values_compound[metric] - minimum = min(values) - maximum = max(values) - avg = np.mean(values) - final_summary["{}_compound".format(metric)] = avg - # print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) - data["{}_compound_min".format(metric)] = minimum - data["{}_compound_max".format(metric)] = maximum - data["{}_compound_avg".format(metric)] = avg + N = np.NAN + if "N" in summary: + N = summary["N"] + Ntotal += N - return final_summary + M = np.NAN + if "M" in summary: + M = summary["M"] + Mtotal += M + + data = { + "layer_id": layer_id, + "layer_type": layer_type, + "N": N, + "M": M, + "slice": slice_id, + "level": "SLICE", + "comment": "Slice level", + } + for metric in metrics: + if metric in summary: + value = summary[metric] + if value is not None: + metrics_values[metric].append(value) + compounds[metric].append(value) + data[metric] = value + + data = { + "layer_id": layer_id, + "layer_type": layer_type, + "N": Ntotal, + "M": Mtotal, + "slice_count": slice_count, + "level": "LAYER", + "comment": "Layer level", + } + # Compute the compound value over the slices + for metric, value in compounds.items(): + count = len(value) + if count == 0: + continue + + compound = np.mean(value) + metrics_values_compound[metric].append(compound) + data[metric] = compound + + data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"} + for metric, metric_name in metrics.items(): + if metric not in metrics_values or len(metrics_values[metric]) == 0: + continue + + values = metrics_values[metric] + minimum = min(values) + maximum = max(values) + avg = np.mean(values) + final_summary[metric] = avg + # print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) + data["{}_min".format(metric)] = minimum + data["{}_max".format(metric)] = maximum + data["{}_avg".format(metric)] = avg + + values = metrics_values_compound[metric] + minimum = min(values) + maximum = max(values) + avg = np.mean(values) + final_summary["{}_compound".format(metric)] = avg + # print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) + data["{}_compound_min".format(metric)] = minimum + data["{}_compound_max".format(metric)] = maximum + data["{}_compound_avg".format(metric)] = avg + + return final_summary -def analyze(model: nn.Module, min_size=50, max_size=0, - alphas: bool = False, lognorms: bool = True, spectralnorms: bool = False, - softranks: bool = False, normalize: bool = False, glorot_fix: bool = False): - """ - Analyze the weight matrices of a model. - :param model: A PyTorch model - :param min_size: The minimum weight matrix size to analyze. - :param max_size: The maximum weight matrix size to analyze (0 = no limit). - :param alphas: Compute the power laws (alpha) of the weight matrices. - Time consuming so disabled by default (use lognorm if you want speed) - :param lognorms: Compute the log norms of the weight matrices. - :param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices. - :param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices. - :param normalize: Normalize or not. - :param glorot_fix: - :return: (a dict of all layers' results, a dict of the summarized info) - """ - names, modules = [], [] - for name, module in model.named_modules(): - if isinstance(module, available_module_types()): - names.append(name) - modules.append(module) - # print('There are {:} layers to be analyzed in this model.'.format(len(modules))) - all_results = OrderedDict() - for index, module in enumerate(modules): - if isinstance(module, nn.Linear): - weights = [module.weight.cpu().detach().numpy()] - else: - weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy()) - results = analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix) - results['id'] = index - results['type'] = type(module) - all_results[index] = results - summary = compute_details(all_results) - return all_results, summary \ No newline at end of file +def analyze( + model: nn.Module, + min_size=50, + max_size=0, + alphas: bool = False, + lognorms: bool = True, + spectralnorms: bool = False, + softranks: bool = False, + normalize: bool = False, + glorot_fix: bool = False, +): + """ + Analyze the weight matrices of a model. + :param model: A PyTorch model + :param min_size: The minimum weight matrix size to analyze. + :param max_size: The maximum weight matrix size to analyze (0 = no limit). + :param alphas: Compute the power laws (alpha) of the weight matrices. + Time consuming so disabled by default (use lognorm if you want speed) + :param lognorms: Compute the log norms of the weight matrices. + :param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices. + :param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices. + :param normalize: Normalize or not. + :param glorot_fix: + :return: (a dict of all layers' results, a dict of the summarized info) + """ + names, modules = [], [] + for name, module in model.named_modules(): + if isinstance(module, available_module_types()): + names.append(name) + modules.append(module) + # print('There are {:} layers to be analyzed in this model.'.format(len(modules))) + all_results = OrderedDict() + for index, module in enumerate(modules): + if isinstance(module, nn.Linear): + weights = [module.weight.cpu().detach().numpy()] + else: + weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy()) + results = analyze_weights( + weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix + ) + results["id"] = index + results["type"] = type(module) + all_results[index] = results + summary = compute_details(all_results) + return all_results, summary