How to use nncf - 10 common examples

def train_epoch_end(config, compression_algo, net, epoch, iteration, epoch_size, lr_scheduler, optimizer,
                    test_data_loader):
    test_freq_in_epochs = max(config.test_interval // epoch_size, 1)
    compression_algo.scheduler.epoch_step(epoch)
    if not isinstance(lr_scheduler, ReduceLROnPlateau):
        lr_scheduler.step(epoch)
    if epoch % test_freq_in_epochs == 0 and iteration != 0:
        if is_on_first_rank(config):
            print_statistics(compression_algo.statistics())
        with torch.no_grad():
            net.eval()
            mAP = test_net(net, config.device, test_data_loader, distributed=config.multiprocessing_distributed)
            if isinstance(lr_scheduler, ReduceLROnPlateau):
                lr_scheduler.step(mAP)
            net.train()
    if epoch > 0 and epoch % config.save_freq == 0 and is_on_first_rank(config):
        print('Saving state, iter:', iteration)
        checkpoint_file_path = osp.join(config.intermediate_checkpoints_path,
                                        "{}_{}.pth".format(config.model, iteration))
        torch.save({
            'state_dict': net.state_dict(),
            'optimizer': optimizer.state_dict(),
            'iter': iteration,
            'scheduler': compression_algo.scheduler.state_dict()
        }, str(checkpoint_file_path))

else:
        dataset = NDG(args.val, args.v_land)

    if dataset.have_landmarks:
        log.info('Use alignment for the train data')
        dataset.transform = t.Compose([Rescale((48, 48)), ToTensor(switch_rb=True)])
    else:
        exit()

    val_loader = DataLoader(dataset, batch_size=args.val_batch_size, num_workers=4, shuffle=False, pin_memory=True)

    model = models_landmarks['landnet']()

    assert args.snapshot is not None
    if args.compr_config:
        config = Config.from_json(args.compr_config)
        compression_algo = create_compression_algorithm(model, config)
        model = compression_algo.model

    log.info('Testing snapshot ' + args.snapshot + ' ...')
    model = load_model_state(model, args.snapshot, args.device, eval_state=True)
    model.eval()
    cudnn.benchmark = True
    model = torch.nn.DataParallel(model, device_ids=[args.device], )

    log.info('Face landmarks model:')
    log.info(model)

    avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)
    log.info('Avg RMSE error: {}'.format(avg_err))
    log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))
    log.info('Failure rate: {}'.format(failures_rate))

dataset = NDG(args.val, args.v_land)

    if dataset.have_landmarks:
        log.info('Use alignment for the train data')
        dataset.transform = t.Compose([Rescale((48, 48)), ToTensor(switch_rb=True)])
    else:
        exit()

    val_loader = DataLoader(dataset, batch_size=args.val_batch_size, num_workers=4, shuffle=False, pin_memory=True)

    model = models_landmarks['landnet']()

    assert args.snapshot is not None
    if args.compr_config:
        config = Config.from_json(args.compr_config)
        compression_algo = create_compression_algorithm(model, config)
        model = compression_algo.model

    log.info('Testing snapshot ' + args.snapshot + ' ...')
    model = load_model_state(model, args.snapshot, args.device, eval_state=True)
    model.eval()
    cudnn.benchmark = True
    model = torch.nn.DataParallel(model, device_ids=[args.device], )

    log.info('Face landmarks model:')
    log.info(model)

    avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)
    log.info('Avg RMSE error: {}'.format(avg_err))
    log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))
    log.info('Failure rate: {}'.format(failures_rate))
    if args.compr_config and "sparsity_level" in compression_algo.statistics():

if config.to_onnx is not None:
        compression_algo.export_model(config.to_onnx)
        print("Saved to", config.to_onnx)
        return

    if config.mode.lower() == 'test':
        print(model)
        model_parameters = filter(lambda p: p.requires_grad, model.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        print("Trainable argument count:{params}".format(params=params))

        model = model.to(config.device)
        loaders, w_class = load_dataset(dataset, config)
        _, val_loader = loaders
        test(model, val_loader, w_class, color_encoding, config)
        print_statistics(compression_algo.statistics())
    elif config.mode.lower() == 'train':
        loaders, w_class = load_dataset(dataset, config)
        train_loader, val_loader = loaders
        if not resuming_checkpoint:
            compression_algo.initialize(train_loader)
        model = \
            train(model, model_without_dp, compression_algo, train_loader, val_loader, w_class, color_encoding, config)
    else:
        # Should never happen...but just in case it does
        raise RuntimeError(
            "\"{0}\" is not a valid choice for execution mode.".format(
                config.mode))

def main(argv):
    parser = get_common_argument_parser()
    arguments = parser.parse_args(args=argv)
    config = Config.from_json(arguments.config)
    config.update_from_args(arguments, parser)
    if config.dist_url == "env://":
        config.update_from_env()

    if config.mode.lower() != 'test':
        if not osp.exists(config.log_dir):
            os.makedirs(config.log_dir)

        config.log_dir = str(config.log_dir)
        configure_paths(config)
        print("Save directory:", config.log_dir)
    else:
        config.log_dir = "/tmp/"

    config.execution_mode = get_execution_mode(config)
    start_worker(main_worker, config)

def main(argv):
    parser = get_argument_parser()
    args = parser.parse_args(args=argv)
    config = Config.from_json(args.config)
    config.update_from_args(args, parser)
    if config.dist_url == "env://":
        config.update_from_env()

    configure_paths(config)
    source_root = Path(__file__).absolute().parents[2]  # nncf root
    create_code_snapshot(source_root, osp.join(config.log_dir, "snapshot.tar.gz"))

    if config.seed is not None:
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    config.execution_mode = get_execution_mode(config)

log_path = './logs/{:%Y_%m_%d_%H_%M}_{}'.format(datetime.datetime.now(), args.snap_prefix)
    writer = SummaryWriter(log_path)

    if not osp.exists(args.snap_folder):
        os.mkdir(args.snap_folder)

    model = models_backbones[args.model](embedding_size=args.embed_size,
                                         num_classes=dataset.get_num_classes(), feature=False)

    set_dropout_fn = model.set_dropout_ratio

    compression_algo = None
    if args.snap_to_resume is not None:
        if args.compr_config:
            config = Config.from_json(args.compr_config)
            compression_algo = create_compression_algorithm(model, config)
            model = compression_algo.model

        log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
        model = load_model_state(model, args.snap_to_resume, args.devices[0], eval_state=False)
        model = torch.nn.DataParallel(model, device_ids=args.devices)
    else:
        model = torch.nn.DataParallel(model, device_ids=args.devices, output_device=args.devices[0])
        model.cuda()
        model.train()
        cudnn.benchmark = True

    if args.to_onnx is not None:
        if args.compr_config:
            compression_algo.export_model(args.to_onnx)
        else:

parser.add_argument('-pp', '--plugin_dir', type=str, default=None, help='Path to a plugin folder')
    parser.add_argument('-c', '--compr_config', help='Path to a file with compression parameters', required=False)
    args = parser.parse_args()

    if args.engine == 'pt':
        assert args.snap is not None, 'To evaluate PyTorch snapshot, please, specify --snap option.'

        if args.compr_config:
            patch_torch_operators()

        with torch.cuda.device(args.devices[0]):
            data, embeddings_fun = load_test_dataset(args)
            model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)

            if args.compr_config:
                config = Config.from_json(args.compr_config)
                compression_algo = create_compression_algorithm(model, config)
                model = compression_algo.model

            model = load_model_state(model, args.snap, args.devices[0])
            evaluate(args, data, model, embeddings_fun, args.val_batch_size, args.dump_embeddings,
                     args.roc_fname, args.snap, True, args.show_failed)

            if args.compr_config and "sparsity_level" in compression_algo.statistics():
                log.info("Sparsity level: {0:.2f}".format(
                    compression_algo.statistics()['sparsity_rate_for_sparsified_modules']))
    else:
        from utils.ie_tools import load_ie_model

        assert args.fr_model is not None, 'To evaluate IE model, please, specify --fr_model option.'
        fr_model = load_ie_model(args.fr_model, 'CPU', args.plugin_dir)
        lm_model = None

def main(argv):
    parser = get_argument_parser()
    args = parser.parse_args(args=argv)
    config = Config.from_json(args.config)
    config.update_from_args(args, parser)
    configure_paths(config)
    source_root = Path(__file__).absolute().parents[2]  # nncf root
    create_code_snapshot(source_root, osp.join(config.log_dir, "snapshot.tar.gz"))

    config.execution_mode = get_execution_mode(config)

    if config.dataset_dir is not None:
        config.train_imgs = config.train_ano = config.test_imgs = config.test_anno = config.dataset_dir
    start_worker(main_worker, config)

- optimizer (``torch.optim``): The stored optimizer state is copied to this
    optimizer instance.
    - compression_algo: The compression scheduler for the saved state
                        to be loaded into

    Returns:
    The ``model``, ``optimizer``, epoch, mean IoU and ``compression_scheduler``, loaded from the
    checkpoint.

    """
    assert os.path.isfile(
        model_path), "The model file \"{0}\" doesn't exist.".format(model_path)

    # Load the stored model parameters to the model instance
    checkpoint = torch.load(model_path, map_location=device_name)
    load_state(model, checkpoint['state_dict'], is_resume=True)
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint['optimizer'])
    epoch = checkpoint['epoch']
    miou = checkpoint['miou']

    if "scheduler" in checkpoint and compression_scheduler is not None:
        compression_scheduler.load_state_dict(checkpoint['scheduler'])

    return model, optimizer, epoch, miou, compression_scheduler