def load_detection_annotations(cachedir, dataset): cachefile = os.path.join(cachedir, 'annots_{}.json'.format(dataset.name)) imagenames = dataset.get_img_names() if is_main_process(): if not os.path.isfile(cachefile): # load annots gt = {} for i, imagename in enumerate(imagenames): _, gt[imagename] = dataset.pull_anno(i) if i % 100 == 0: logger.info('Reading annotation for {:d}/{:d}'.format( i + 1, len(imagenames))) # save logger.info('Saving cached annotations to {:s}'.format(cachefile)) pathlib.Path(cachedir).mkdir(parents=True, exist_ok=True) with open(cachefile, 'w', encoding='utf8') as f: json.dump(gt, f) if is_dist_avail_and_initialized(): dist.barrier() with open(cachefile, 'r', encoding='utf8') as f: gt = json.load(f) return gt, imagenames
def train_epoch_staged(train_loader, batch_multiplier, model, criterion, criterion_fn, optimizer, optimizer_scheduler: PolyLRDropScheduler, kd_loss_calculator: KDLossCalculator, compression_ctrl, epoch, config): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() kd_losses_meter = AverageMeter() criterion_losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() compression_scheduler = compression_ctrl.scheduler # switch to train mode model.train() end = time.time() for i, (input_, target) in enumerate(train_loader): compression_scheduler.step() # measure data loading time data_time.update(time.time() - end) input_ = input_.to(config.device) target = target.to(config.device) output = model(input_) criterion_loss = criterion_fn(output, target, criterion) if isinstance(output, InceptionOutputs): output = output.logits # compute KD loss kd_loss = kd_loss_calculator.loss(input_, output) loss = criterion_loss + kd_loss # measure accuracy and record loss acc1, acc5 = accuracy(output, target, topk=(1, 5)) losses.update(loss.item(), input_.size(0)) comp_loss_val = kd_loss.item() kd_losses_meter.update(comp_loss_val, input_.size(0)) criterion_losses.update(criterion_loss.item(), input_.size(0)) top1.update(acc1, input_.size(0)) top1.update(acc1, input_.size(0)) top5.update(acc5, input_.size(0)) # compute gradient and do SGD step if i % batch_multiplier == 0: optimizer.zero_grad() loss.backward() optimizer.step() else: loss.backward() optimizer_scheduler.step(float(i) / len(train_loader)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % config.print_freq == 0: logger.info('{rank}: ' 'Epoch: [{0}][{1}/{2}] ' 'Lr: {3:.3} ' 'Wd: {4:.3} ' 'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data: {data_time.val:.3f} ({data_time.avg:.3f}) ' 'CE_loss: {ce_loss.val:.4f} ({ce_loss.avg:.4f}) ' 'KD_loss: {kd_loss.val:.4f} ({kd_loss.avg:.4f}) ' 'Loss: {loss.val:.4f} ({loss.avg:.4f}) ' 'Acc@1: {top1.val:.3f} ({top1.avg:.3f}) ' 'Acc@5: {top5.val:.3f} ({top5.avg:.3f})'.format( epoch, i, len(train_loader), get_lr(optimizer), get_wd(optimizer), batch_time=batch_time, data_time=data_time, ce_loss=criterion_losses, kd_loss=kd_losses_meter, loss=losses, top1=top1, top5=top5, rank='{}:'.format(config.rank) if config.multiprocessing_distributed else '')) if is_main_process(): global_step = len(train_loader) * epoch config.tb.add_scalar("train/learning_rate", get_lr(optimizer), i + global_step) config.tb.add_scalar("train/criterion_loss", criterion_losses.avg, i + global_step) config.tb.add_scalar("train/kd_loss", kd_losses_meter.avg, i + global_step) config.tb.add_scalar("train/loss", losses.avg, i + global_step) config.tb.add_scalar("train/top1", top1.avg, i + global_step) config.tb.add_scalar("train/top5", top5.avg, i + global_step) statistics = compression_ctrl.statistics( quickly_collected_only=True) for stat_name, stat_value in prepare_for_tensorboard( statistics).items(): config.tb.add_scalar('train/statistics/{}'.format(stat_name), stat_value, i + global_step)
def train_staged(config, compression_ctrl, model, criterion, criterion_fn, optimizer_scheduler, model_name, optimizer, train_loader, train_sampler, val_loader, kd_loss_calculator, batch_multiplier, best_acc1=0): best_compression_stage = CompressionStage.UNCOMPRESSED for epoch in range(config.start_epoch, config.epochs): # update compression scheduler state at the start of the epoch compression_ctrl.scheduler.epoch_step() if config.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_epoch_staged(train_loader, batch_multiplier, model, criterion, criterion_fn, optimizer, optimizer_scheduler, kd_loss_calculator, compression_ctrl, epoch, config) # compute compression algo statistics statistics = compression_ctrl.statistics() acc1 = best_acc1 if epoch % config.test_every_n_epochs == 0: # evaluate on validation set # pylint: disable=E1123 acc1, _, _ = validate(val_loader, model, criterion, config, epoch=epoch) compression_stage = compression_ctrl.compression_stage() # remember best acc@1, considering compression stage. If current acc@1 less then the best acc@1, checkpoint # still can be best if current compression stage is larger than the best one. Compression stages in ascending # order: UNCOMPRESSED, PARTIALLY_COMPRESSED, FULLY_COMPRESSED. is_best_by_accuracy = acc1 > best_acc1 and compression_stage == best_compression_stage is_best = is_best_by_accuracy or compression_stage > best_compression_stage best_acc1 = max(acc1, best_acc1) best_compression_stage = max(compression_stage, best_compression_stage) # statistics (e.g. portion of the enabled quantizers) is related to the finished epoch, # hence printing should happen before epoch_step, which may inform about state of the next epoch (e.g. next # portion of enabled quantizers) if is_main_process(): logger.info(statistics.to_str()) optimizer_scheduler.epoch_step() if is_main_process(): checkpoint_path = osp.join(config.checkpoint_save_dir, get_name(config) + '_last.pth') checkpoint = { 'epoch': epoch + 1, 'arch': model_name, MODEL_STATE_ATTR: model.state_dict(), COMPRESSION_STATE_ATTR: compression_ctrl.get_compression_state(), 'original_model_state_dict': kd_loss_calculator.original_model.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), 'optimizer_scheduler': optimizer_scheduler.state_dict() } torch.save(checkpoint, checkpoint_path) make_additional_checkpoints(checkpoint_path, is_best, epoch + 1, config) for key, value in prepare_for_tensorboard(statistics).items(): config.mlflow.safe_call( 'log_metric', 'compression/statistics/{0}'.format(key), value, epoch) config.tb.add_scalar("compression/statistics/{0}".format(key), value, len(train_loader) * epoch)
def staged_quantization_main_worker(current_gpu, config): configure_device(current_gpu, config) config.mlflow = SafeMLFLow(config) if is_main_process(): configure_logging(logger, config) print_args(config) set_seed(config) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss() criterion = criterion.to(config.device) model_name = config['model'] is_inception = 'inception' in model_name train_criterion_fn = inception_criterion_fn if is_inception else default_criterion_fn train_loader = train_sampler = val_loader = None resuming_checkpoint_path = config.resuming_checkpoint_path nncf_config = config.nncf_config pretrained = is_pretrained_model_requested(config) is_export_only = 'export' in config.mode and ( 'train' not in config.mode and 'test' not in config.mode) if is_export_only: assert pretrained or (resuming_checkpoint_path is not None) else: # Data loading code train_dataset, val_dataset = create_datasets(config) train_loader, train_sampler, val_loader, init_loader = create_data_loaders( config, train_dataset, val_dataset) def autoq_eval_fn(model, eval_loader): _, top5, _ = validate(eval_loader, model, criterion, config) return top5 nncf_config = register_default_init_args( nncf_config, init_loader, criterion=criterion, criterion_fn=train_criterion_fn, autoq_eval_fn=autoq_eval_fn, val_loader=val_loader, device=config.device) # create model model_name = config['model'] model = load_model(model_name, pretrained=pretrained, num_classes=config.get('num_classes', 1000), model_params=config.get('model_params'), weights_path=config.get('weights')) original_model = copy.deepcopy(model) model.to(config.device) resuming_checkpoint = None if resuming_checkpoint_path is not None: resuming_checkpoint = load_resuming_checkpoint( resuming_checkpoint_path) model_state_dict, compression_state = extract_model_and_compression_states( resuming_checkpoint) compression_ctrl, model = create_compressed_model(model, nncf_config, compression_state) if model_state_dict is not None: load_state(model, model_state_dict, is_resume=True) if not isinstance(compression_ctrl, (BinarizationController, QuantizationController)): raise RuntimeError( "The stage quantization sample worker may only be run with the binarization and quantization algorithms!" ) model, _ = prepare_model_for_execution(model, config) original_model.to(config.device) if config.distributed: compression_ctrl.distributed() params_to_optimize = model.parameters() compression_config = config['compression'] quantization_config = compression_config if isinstance( compression_config, dict) else compression_config[0] optimizer = get_quantization_optimizer(params_to_optimize, quantization_config) optimizer_scheduler = PolyLRDropScheduler(optimizer, quantization_config) kd_loss_calculator = KDLossCalculator(original_model) best_acc1 = 0 # optionally resume from a checkpoint if resuming_checkpoint is not None and config.to_onnx is None: config.start_epoch = resuming_checkpoint['epoch'] best_acc1 = resuming_checkpoint['best_acc1'] kd_loss_calculator.original_model.load_state_dict( resuming_checkpoint['original_model_state_dict']) if 'train' in config.mode: optimizer.load_state_dict(resuming_checkpoint['optimizer']) optimizer_scheduler.load_state_dict( resuming_checkpoint['optimizer_scheduler']) logger.info( "=> loaded checkpoint '{}' (epoch: {}, best_acc1: {:.3f})". format(resuming_checkpoint_path, resuming_checkpoint['epoch'], best_acc1)) else: logger.info( "=> loaded checkpoint '{}'".format(resuming_checkpoint_path)) log_common_mlflow_params(config) if is_export_only: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx)) return if config.execution_mode != ExecutionMode.CPU_ONLY: cudnn.benchmark = True if is_main_process(): statistics = compression_ctrl.statistics() logger.info(statistics.to_str()) if 'train' in config.mode: batch_multiplier = (quantization_config.get("params", {})).get( "batch_multiplier", 1) train_staged(config, compression_ctrl, model, criterion, train_criterion_fn, optimizer_scheduler, model_name, optimizer, train_loader, train_sampler, val_loader, kd_loss_calculator, batch_multiplier, best_acc1) if 'test' in config.mode: validate(val_loader, model, criterion, config) if 'export' in config.mode: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx))
def train(net, compression_ctrl, train_data_loader, test_data_loader, criterion, optimizer, config, lr_scheduler): net.train() loc_loss = 0 conf_loss = 0 epoch_size = len(train_data_loader) logger.info('Training {} on {} dataset...'.format( config.model, train_data_loader.dataset.name)) best_mAp = 0 best_compression_stage = CompressionStage.UNCOMPRESSED test_freq_in_epochs = config.test_interval if config.test_interval is None: test_freq_in_epochs = 1 max_epochs = config['epochs'] for epoch in range(config.start_epoch, max_epochs): compression_ctrl.scheduler.epoch_step(epoch) train_epoch(compression_ctrl, net, config, train_data_loader, criterion, optimizer, epoch_size, epoch, loc_loss, conf_loss) if is_main_process(): logger.info(compression_ctrl.statistics().to_str()) compression_stage = compression_ctrl.compression_stage() is_best = False if (epoch + 1) % test_freq_in_epochs == 0: with torch.no_grad(): net.eval() mAP = test_net(net, config.device, test_data_loader, distributed=config.multiprocessing_distributed) is_best_by_mAP = mAP > best_mAp and compression_stage == best_compression_stage is_best = is_best_by_mAP or compression_stage > best_compression_stage if is_best: best_mAp = mAP best_compression_stage = max(compression_stage, best_compression_stage) if isinstance(lr_scheduler, ReduceLROnPlateau): lr_scheduler.step(mAP) net.train() if is_on_first_rank(config): logger.info('Saving state, epoch: {}'.format(epoch)) checkpoint_file_path = osp.join( config.checkpoint_save_dir, "{}_last.pth".format(get_name(config))) torch.save( { MODEL_STATE_ATTR: net.state_dict(), COMPRESSION_STATE_ATTR: compression_ctrl.get_compression_state(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, }, str(checkpoint_file_path)) make_additional_checkpoints(checkpoint_file_path, is_best=is_best, epoch=epoch + 1, config=config) # Learning rate scheduling should be applied after optimizer’s update if not isinstance(lr_scheduler, ReduceLROnPlateau): lr_scheduler.step(epoch) compression_ctrl.scheduler.epoch_step(epoch) if is_main_process(): statistics = compression_ctrl.statistics() logger.info(statistics.to_str()) compression_stage = compression_ctrl.compression_stage() is_best = False if (epoch + 1) % test_freq_in_epochs == 0: with torch.no_grad(): net.eval() mAP = test_net(net, config.device, test_data_loader, distributed=config.multiprocessing_distributed) is_best_by_mAP = mAP > best_mAp and compression_stage == best_compression_stage is_best = is_best_by_mAP or compression_stage > best_compression_stage if is_best: best_mAp = mAP best_compression_stage = max(compression_stage, best_compression_stage) if isinstance(lr_scheduler, ReduceLROnPlateau): lr_scheduler.step(mAP) net.train() if is_on_first_rank(config): logger.info('Saving state, epoch: {}'.format(epoch)) checkpoint_file_path = osp.join( config.checkpoint_save_dir, "{}_last.pth".format(get_name(config))) torch.save( { MODEL_STATE_ATTR: net.state_dict(), COMPRESSION_STATE_ATTR: compression_ctrl.get_compression_state(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, }, str(checkpoint_file_path)) make_additional_checkpoints(checkpoint_file_path, is_best=is_best, epoch=epoch + 1, config=config) # Learning rate scheduling should be applied after optimizer’s update if not isinstance(lr_scheduler, ReduceLROnPlateau): lr_scheduler.step(epoch) if config.metrics_dump is not None: write_metrics(best_mAp, config.metrics_dump)
def main_worker(current_gpu, config): ################################# # Setup experiment environment ################################# configure_device(current_gpu, config) config.mlflow = SafeMLFLow(config) if is_on_first_rank(config): configure_logging(logger, config) print_args(config) set_seed(config) config.start_iter = 0 nncf_config = config.nncf_config ########################## # Prepare metrics log file ########################## if config.metrics_dump is not None: write_metrics(0, config.metrics_dump) ########################### # Criterion ########################### criterion = MultiBoxLoss(config, config['num_classes'], overlap_thresh=0.5, prior_for_matching=True, bkg_label=0, neg_mining=True, neg_pos=3, neg_overlap=0.5, encode_target=False, device=config.device) train_data_loader = test_data_loader = None resuming_checkpoint_path = config.resuming_checkpoint_path ########################### # Prepare data ########################### pretrained = is_pretrained_model_requested(config) is_export_only = 'export' in config.mode and ( 'train' not in config.mode and 'test' not in config.mode) if is_export_only: assert pretrained or (resuming_checkpoint_path is not None) else: test_data_loader, train_data_loader, init_data_loader = create_dataloaders( config) def criterion_fn(model_outputs, target, criterion): loss_l, loss_c = criterion(model_outputs, target) return loss_l + loss_c def autoq_test_fn(model, eval_loader): # RL is maximization, change the loss polarity return -1 * test_net(model, config.device, eval_loader, distributed=config.distributed, loss_inference=True, criterion=criterion) def model_eval_fn(model): model.eval() mAP = test_net(model, config.device, test_data_loader, distributed=config.distributed, criterion=criterion) return mAP nncf_config = register_default_init_args(nncf_config, init_data_loader, criterion=criterion, criterion_fn=criterion_fn, autoq_eval_fn=autoq_test_fn, val_loader=test_data_loader, model_eval_fn=model_eval_fn, device=config.device) ################## # Prepare model ################## resuming_checkpoint_path = config.resuming_checkpoint_path resuming_checkpoint = None if resuming_checkpoint_path is not None: resuming_checkpoint = load_resuming_checkpoint( resuming_checkpoint_path) compression_ctrl, net = create_model(config, resuming_checkpoint) if config.distributed: config.batch_size //= config.ngpus_per_node config.workers //= config.ngpus_per_node compression_ctrl.distributed() ########################### # Optimizer ########################### params_to_optimize = get_parameter_groups(net, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) ################################# # Load additional checkpoint data ################################# if resuming_checkpoint_path is not None and 'train' in config.mode: optimizer.load_state_dict( resuming_checkpoint.get('optimizer', optimizer.state_dict())) config.start_epoch = resuming_checkpoint.get('epoch', 0) + 1 log_common_mlflow_params(config) if is_export_only: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx)) return if is_main_process(): statistics = compression_ctrl.statistics() logger.info(statistics.to_str()) if 'train' in config.mode and is_accuracy_aware_training(config): # validation function that returns the target metric value # pylint: disable=E1123 def validate_fn(model, epoch): model.eval() mAP = test_net(model, config.device, test_data_loader, distributed=config.distributed) model.train() return mAP # training function that trains the model for one epoch (full training dataset pass) # it is assumed that all the NNCF-related methods are properly called inside of # this function (like e.g. the step and epoch_step methods of the compression scheduler) def train_epoch_fn(compression_ctrl, model, epoch, optimizer, **kwargs): loc_loss = 0 conf_loss = 0 epoch_size = len(train_data_loader) train_epoch(compression_ctrl, model, config, train_data_loader, criterion, optimizer, epoch_size, epoch, loc_loss, conf_loss) # function that initializes optimizers & lr schedulers to start training def configure_optimizers_fn(): params_to_optimize = get_parameter_groups(net, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) return optimizer, lr_scheduler acc_aware_training_loop = create_accuracy_aware_training_loop( nncf_config, compression_ctrl) net = acc_aware_training_loop.run( net, train_epoch_fn=train_epoch_fn, validate_fn=validate_fn, configure_optimizers_fn=configure_optimizers_fn, tensorboard_writer=config.tb, log_dir=config.log_dir) elif 'train' in config.mode: train(net, compression_ctrl, train_data_loader, test_data_loader, criterion, optimizer, config, lr_scheduler) if 'test' in config.mode: with torch.no_grad(): net.eval() if config['ssd_params'].get('loss_inference', False): model_loss = test_net(net, config.device, test_data_loader, distributed=config.distributed, loss_inference=True, criterion=criterion) logger.info("Final model loss: {:.3f}".format(model_loss)) else: mAp = test_net(net, config.device, test_data_loader, distributed=config.distributed) if config.metrics_dump is not None: write_metrics(mAp, config.metrics_dump) if 'export' in config.mode: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx))
def _is_enabled(self): return self.is_suitable_mode and is_main_process()
def main_worker(current_gpu, config): configure_device(current_gpu, config) config.mlflow = SafeMLFLow(config) if is_main_process(): configure_logging(logger, config) print_args(config) set_seed(config) logger.info(config) dataset = get_dataset(config.dataset) color_encoding = dataset.color_encoding num_classes = len(color_encoding) if config.metrics_dump is not None: write_metrics(0, config.metrics_dump) train_loader = val_loader = criterion = None resuming_checkpoint_path = config.resuming_checkpoint_path nncf_config = config.nncf_config pretrained = is_pretrained_model_requested(config) def criterion_fn(model_outputs, target, criterion_): labels, loss_outputs, _ = \ loss_funcs.do_model_specific_postprocessing(config.model, target, model_outputs) return criterion_(loss_outputs, labels) is_export_only = 'export' in config.mode and ( 'train' not in config.mode and 'test' not in config.mode) if is_export_only: assert pretrained or (resuming_checkpoint_path is not None) else: loaders, w_class = load_dataset(dataset, config) train_loader, val_loader, init_loader = loaders criterion = get_criterion(w_class, config) def autoq_test_fn(model, eval_loader): return test(model, eval_loader, criterion, color_encoding, config) model_eval_fn = functools.partial(autoq_test_fn, eval_loader=val_loader) nncf_config = register_default_init_args(nncf_config, init_loader, criterion=criterion, criterion_fn=criterion_fn, autoq_eval_fn=autoq_test_fn, val_loader=val_loader, model_eval_fn=model_eval_fn, device=config.device) model = load_model(config.model, pretrained=pretrained, num_classes=num_classes, model_params=config.get('model_params', {}), weights_path=config.get('weights')) model.to(config.device) resuming_checkpoint = None if resuming_checkpoint_path is not None: resuming_checkpoint = load_resuming_checkpoint( resuming_checkpoint_path) model_state_dict, compression_state = extract_model_and_compression_states( resuming_checkpoint) compression_ctrl, model = create_compressed_model(model, nncf_config, compression_state) if model_state_dict is not None: load_state(model, model_state_dict, is_resume=True) model, model_without_dp = prepare_model_for_execution(model, config) if config.distributed: compression_ctrl.distributed() log_common_mlflow_params(config) if is_export_only: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx)) return if is_main_process(): statistics = compression_ctrl.statistics() logger.info(statistics.to_str()) if is_accuracy_aware_training(config) and 'train' in config.mode: def validate_fn(model, epoch): return test(model, val_loader, criterion, color_encoding, config) # training function that trains the model for one epoch (full training dataset pass) # it is assumed that all the NNCF-related methods are properly called inside of # this function (like e.g. the step and epoch_step methods of the compression scheduler) def train_epoch_fn(compression_ctrl, model, optimizer, **kwargs): ignore_index = None ignore_unlabeled = config.get("ignore_unlabeled", True) if ignore_unlabeled and ('unlabeled' in color_encoding): ignore_index = list(color_encoding).index('unlabeled') metric = IoU(len(color_encoding), ignore_index=ignore_index) train_obj = Train(model, train_loader, optimizer, criterion, compression_ctrl, metric, config.device, config.model) train_obj.run_epoch(config.print_step) # function that initializes optimizers & lr schedulers to start training def configure_optimizers_fn(): optim_config = config.get('optimizer', {}) optim_params = optim_config.get('optimizer_params', {}) lr = optim_params.get("lr", 1e-4) params_to_optimize = get_params_to_optimize( model_without_dp, lr * 10, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) return optimizer, lr_scheduler acc_aware_training_loop = create_accuracy_aware_training_loop( config, compression_ctrl) model = acc_aware_training_loop.run( model, train_epoch_fn=train_epoch_fn, validate_fn=validate_fn, configure_optimizers_fn=configure_optimizers_fn, tensorboard_writer=config.tb, log_dir=config.log_dir) elif 'train' in config.mode: train(model, model_without_dp, compression_ctrl, train_loader, val_loader, criterion, color_encoding, config, resuming_checkpoint) if 'test' in config.mode: logger.info(model) model_parameters = filter(lambda p: p.requires_grad, model.parameters()) params = sum([np.prod(p.size()) for p in model_parameters]) logger.info("Trainable argument count:{params}".format(params=params)) model = model.to(config.device) test(model, val_loader, criterion, color_encoding, config) if 'export' in config.mode: compression_ctrl.export_model(config.to_onnx) logger.info("Saved to {}".format(config.to_onnx))
def train(model, model_without_dp, compression_ctrl, train_loader, val_loader, criterion, class_encoding, config, resuming_checkpoint): logger.info("\nTraining...\n") # Check if the network architecture is correct logger.info(model) optim_config = config.get('optimizer', {}) optim_params = optim_config.get('optimizer_params', {}) lr = optim_params.get("lr", 1e-4) params_to_optimize = get_params_to_optimize(model_without_dp, lr * 10, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) # Evaluation metric ignore_index = None ignore_unlabeled = config.get("ignore_unlabeled", True) if ignore_unlabeled and ('unlabeled' in class_encoding): ignore_index = list(class_encoding).index('unlabeled') metric = IoU(len(class_encoding), ignore_index=ignore_index) best_miou = -1 best_compression_stage = CompressionStage.UNCOMPRESSED # Optionally resume from a checkpoint if resuming_checkpoint is not None: if optimizer is not None: optimizer.load_state_dict(resuming_checkpoint['optimizer']) start_epoch = resuming_checkpoint['epoch'] best_miou = resuming_checkpoint['miou'] logger.info("Resuming from model: Start epoch = {0} " "| Best mean IoU = {1:.4f}".format(start_epoch, best_miou)) config.start_epoch = start_epoch # Start Training train_obj = Train(model, train_loader, optimizer, criterion, compression_ctrl, metric, config.device, config.model) val_obj = Test(model, val_loader, criterion, metric, config.device, config.model) for epoch in range(config.start_epoch, config.epochs): compression_ctrl.scheduler.epoch_step() logger.info(">>>> [Epoch: {0:d}] Training".format(epoch)) if config.distributed: train_loader.sampler.set_epoch(epoch) epoch_loss, (iou, miou) = train_obj.run_epoch(config.print_step) if not isinstance(lr_scheduler, ReduceLROnPlateau): # Learning rate scheduling should be applied after optimizer’s update lr_scheduler.step(epoch) logger.info( ">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}". format(epoch, epoch_loss, miou)) if is_main_process(): config.tb.add_scalar("train/loss", epoch_loss, epoch) config.tb.add_scalar("train/mIoU", miou, epoch) config.tb.add_scalar("train/learning_rate", optimizer.param_groups[0]['lr'], epoch) config.tb.add_scalar("train/compression_loss", compression_ctrl.loss(), epoch) statistics = compression_ctrl.statistics( quickly_collected_only=True) for key, value in prepare_for_tensorboard(statistics).items(): config.tb.add_scalar("compression/statistics/{0}".format(key), value, epoch) if (epoch + 1) % config.save_freq == 0 or epoch + 1 == config.epochs: logger.info(">>>> [Epoch: {0:d}] Validation".format(epoch)) loss, (iou, miou) = val_obj.run_epoch(config.print_step) logger.info( ">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}". format(epoch, loss, miou)) if is_main_process(): config.tb.add_scalar("val/mIoU", miou, epoch) config.tb.add_scalar("val/loss", loss, epoch) for i, (key, class_iou) in enumerate(zip(class_encoding.keys(), iou)): config.tb.add_scalar( "{}/mIoU_Cls{}_{}".format(config.dataset, i, key), class_iou, epoch) compression_stage = compression_ctrl.compression_stage() is_best_by_miou = miou > best_miou and compression_stage == best_compression_stage is_best = is_best_by_miou or compression_stage > best_compression_stage if is_best: best_miou = miou best_compression_stage = max(compression_stage, best_compression_stage) if config.metrics_dump is not None: write_metrics(best_miou, config.metrics_dump) if isinstance(lr_scheduler, ReduceLROnPlateau): # Learning rate scheduling should be applied after optimizer’s update lr_scheduler.step(best_miou) # Print per class IoU on last epoch or if best iou if epoch + 1 == config.epochs or is_best: for key, class_iou in zip(class_encoding.keys(), iou): logger.info("{0}: {1:.4f}".format(key, class_iou)) # Save the model if it's the best thus far if is_main_process(): checkpoint_path = save_checkpoint(model, compression_ctrl, optimizer, epoch, best_miou, config) make_additional_checkpoints(checkpoint_path, is_best, epoch, config) statistics = compression_ctrl.statistics() logger.info(statistics.to_str()) return model