def training(local_rank, config):
    rank = idist.get_rank()
    manual_seed((config['seed'] + rank))
    device = idist.device()
    logger = setup_logger(name='CIFAR10-Training', distributed_rank=local_rank)
    log_basic_info(logger, config)
    output_path = config['output_path']
    if (rank == 0):
        if (config['stop_iteration'] is None):
            now = datetime.now().strftime('%Y%m%d-%H%M%S')
        else:
            now = 'stop-on-{}'.format(config['stop_iteration'])
        folder_name = '{}_backend-{}-{}_{}'.format(config['model'], idist.backend(), idist.get_world_size(), now)
        output_path = (Path(output_path) / folder_name)
        if (not output_path.exists()):
            output_path.mkdir(parents=True)
        config['output_path'] = output_path.as_posix()
        logger.info('Output path: {}'.format(config['output_path']))
        if ('cuda' in device.type):
            config['cuda device name'] = torch.cuda.get_device_name(local_rank)
        if config['with_trains']:
            from trains import Task
            task = Task.init('CIFAR10-Training', task_name=output_path.stem)
            task.connect_configuration(config)
            hyper_params = ['model', 'batch_size', 'momentum', 'weight_decay', 'num_epochs', 'learning_rate', 'num_warmup_epochs']
            task.connect({k: config[k] for k in hyper_params})
    (train_loader, test_loader) = get_dataflow(config)
    config['num_iters_per_epoch'] = len(train_loader)
    (model, optimizer, criterion, lr_scheduler) = initialize(config)
    trainer = create_trainer(model, optimizer, criterion, lr_scheduler, train_loader.sampler, config, logger)
    metrics = {'accuracy': Accuracy(), 'loss': Loss(criterion)}
    evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)
    train_evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)

    def run_validation(engine):
        epoch = trainer.state.epoch
        state = train_evaluator.run(train_loader)
        log_metrics(logger, epoch, state.times['COMPLETED'], 'Train', state.metrics)
        state = evaluator.run(test_loader)
        log_metrics(logger, epoch, state.times['COMPLETED'], 'Test', state.metrics)
    trainer.add_event_handler((Events.EPOCH_COMPLETED(every=config['validate_every']) | Events.COMPLETED), run_validation)
    if (rank == 0):
        evaluators = {'training': train_evaluator, 'test': evaluator}
        tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)
    common.gen_save_best_models_by_val_score(save_handler=get_save_handler(config), evaluator=evaluator, models={'model': model}, metric_name='accuracy', n_saved=3, trainer=trainer, tag='test')
    if (config['stop_iteration'] is not None):

        @trainer.on(Events.ITERATION_STARTED(once=config['stop_iteration']))
        def _():
            logger.info('Stop training on {} iteration'.format(trainer.state.iteration))
            trainer.terminate()
    try:
        trainer.run(train_loader, max_epochs=config['num_epochs'])
    except Exception as e:
        import traceback
        print(traceback.format_exc())
    if (rank == 0):
        tb_logger.close()
def training(local_rank, config):
    rank = idist.get_rank()
    manual_seed((config['seed'] + rank))
    device = idist.device()
    logger = setup_logger(name='CIFAR10-Training', distributed_rank=local_rank)
    log_basic_info(logger, config)
    output_path = config['output_path']
    if (rank == 0):
        if (config['stop_iteration'] is None):
            now = datetime.now().strftime('%Y%m%d-%H%M%S')
        else:
            now = 'stop-on-{}'.format(config['stop_iteration'])
        folder_name = '{}_backend-{}-{}_{}'.format(config['model'], idist.backend(), idist.get_world_size(), now)
        output_path = (Path(output_path) / folder_name)
        if (not output_path.exists()):
            output_path.mkdir(parents=True)
        config['output_path'] = output_path.as_posix()
        logger.info('Output path: {}'.format(config['output_path']))
        if ('cuda' in device.type):
            config['cuda device name'] = torch.cuda.get_device_name(local_rank)
        if config['with_trains']:
            from trains import Task
            task = Task.init('CIFAR10-Training', task_name=output_path.stem)
            task.connect_configuration(config)
            hyper_params = ['model', 'batch_size', 'momentum', 'weight_decay', 'num_epochs', 'learning_rate', 'num_warmup_epochs']
            task.connect({k: config[k] for k in hyper_params})
    (train_loader, test_loader) = get_dataflow(config)
    config['num_iters_per_epoch'] = len(train_loader)
    (model, optimizer, criterion, lr_scheduler) = initialize(config)
    trainer = create_trainer(model, optimizer, criterion, lr_scheduler, train_loader.sampler, config, logger)
    metrics = {'accuracy': Accuracy(), 'loss': Loss(criterion)}
    evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)
    train_evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)

    def run_validation(engine):
        epoch = trainer.state.epoch
        state = train_evaluator.run(train_loader)
        log_metrics(logger, epoch, state.times['COMPLETED'], 'Train', state.metrics)
        state = evaluator.run(test_loader)
        log_metrics(logger, epoch, state.times['COMPLETED'], 'Test', state.metrics)
    trainer.add_event_handler((Events.EPOCH_COMPLETED(every=config['validate_every']) | Events.COMPLETED), run_validation)
    if (rank == 0):
        evaluators = {'training': train_evaluator, 'test': evaluator}
        tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)
    common.gen_save_best_models_by_val_score(save_handler=get_save_handler(config), evaluator=evaluator, models={'model': model}, metric_name='accuracy', n_saved=3, trainer=trainer, tag='test')
    if (config['stop_iteration'] is not None):

        @trainer.on(Events.ITERATION_STARTED(once=config['stop_iteration']))
        def _():
            logger.info('Stop training on {} iteration'.format(trainer.state.iteration))
            trainer.terminate()
    try:
        trainer.run(train_loader, max_epochs=config['num_epochs'])
    except Exception as e:
        import traceback
        print(traceback.format_exc())
    if (rank == 0):
        tb_logger.close()