def main()

in tools/prune.py [0:0]


def main():
    args = parse_args()

    if args.model_type is not None and args.config is None:
        assert args.model_type in CONFIG_TEMPLATE_ZOO, 'model_type must be in [%s]' % (
            ', '.join(CONFIG_TEMPLATE_ZOO.keys()))
        print('model_type=%s, config file will be replaced by %s' %
              (args.model_type, CONFIG_TEMPLATE_ZOO[args.model_type]))
        args.config = CONFIG_TEMPLATE_ZOO[args.model_type]

    if args.config.startswith('http'):

        r = requests.get(args.config)
        # download config in current dir
        tpath = args.config.split('/')[-1]
        while not osp.exists(tpath):
            try:
                with open(tpath, 'wb') as code:
                    code.write(r.content)
            except:
                pass

        args.config = tpath

    cfg = mmcv_config_fromfile(args.config)

    if args.user_config_params is not None:
        assert args.model_type is not None, 'model_type must be setted'
        # rebuild config by user config params
        cfg = rebuild_config(cfg, args.user_config_params)

    # check oss_config and init oss io
    if cfg.get('oss_io_config', None) is not None:
        io.access_oss(**cfg.oss_io_config)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # update configs according to CLI args
    if args.work_dir is not None:
        cfg.work_dir = args.work_dir

    # if `work_dir` is oss path, redirect `work_dir` to local path, add `oss_work_dir` point to oss path,
    # and use osssync hook to upload log and ckpt in work_dir to oss_work_dir
    if cfg.work_dir.startswith('oss://'):
        cfg.oss_work_dir = cfg.work_dir
        cfg.work_dir = osp.join('work_dirs',
                                cfg.work_dir.replace('oss://', ''))
    else:
        cfg.oss_work_dir = None

    # create work_dir
    if not io.exists(cfg.work_dir):
        io.makedirs(cfg.work_dir)

    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    if args.model_type is not None:
        model_name = args.model_type
    else:
        model_name = cfg.model.type

    dummy_input = torch.randn([1, 3, cfg.img_scale[0],
                               cfg.img_scale[1]]).to('cuda')
    prune_layer_names = get_prune_layer(cfg.model.type)

    # set random seeds
    if args.seed is not None:
        logger.info('Set random seed to {}, deterministic: {}'.format(
            args.seed, args.deterministic))
        set_random_seed(args.seed, deterministic=args.deterministic)
    cfg.seed = args.seed

    cfg.gpus = args.gpus

    model = build_model(cfg.model)

    checkpoint = load_checkpoint(model, args.checkpoint)

    if cfg.model.type == 'YOLOX_EDGE' or 'YOLOX':
        prune_optimizer = build_yolo_optimizer(model, cfg.optimizer)
    else:
        prune_optimizer = build_optimizer(model, cfg.optimizer)

    model_layers = []
    for name, module in model.named_modules():
        model_layers.append(name)
    print(model_layers)

    meta = None
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    distributed = False
    # build dataloader
    shuffle = cfg.data.train.pop('shuffle', True)
    print(f'data shuffle: {shuffle}')

    datasets = [build_dataset(cfg.data.train)]
    data_loaders = [
        build_dataloader(
            ds,
            cfg.data.imgs_per_gpu,
            cfg.data.workers_per_gpu,
            cfg.gpus,
            dist=distributed,
            shuffle=shuffle,
            replace=getattr(cfg.data, 'sampling_replace', False),
            seed=cfg.seed,
            drop_last=getattr(cfg.data, 'drop_last', False),
            reuse_worker_cache=cfg.data.get('reuse_worker_cache', False))
        for ds in datasets
    ]

    iters = datasets[0].num_samples * cfg.total_epochs / \
        (get_num_gpu_per_node() * cfg.data.imgs_per_gpu)
    assert int(
        iters
    ) >= 1200, 'pruner need iters larger than 1200, please increase epoch'

    pruner_config = [{
        'sparsity': 0.5,
        'frequency': 200,
        'start_iter': 0,
        'end_iter': 1200,
        'op_names': prune_layer_names
    }]

    pruner = load_pruner(model, pruner_config, prune_optimizer,
                         args.pruning_class, args.pruning_algorithm)

    train_model(
        model,
        data_loaders,
        cfg,
        distributed=distributed,
        timestamp=timestamp,
        meta=meta)

    # export mask model
    mask_path = osp.join(cfg.work_dir, model_name + '_mask.pth')
    model_path = osp.join(cfg.work_dir, model_name + '_model.pth')
    model.head.decode_in_inference = False
    pruner.export_model(model_path, mask_path)

    model = build_model(cfg.model)
    model.load_state_dict(torch.load(model_path))
    model.head.decode_in_inference = False
    model.cuda()
    model.eval()

    # for pai-nni >= 2.6.1 the usr_defined_module should be set and .clone is not supported
    m_speedup = ModelSpeedup(
        model, dummy_input, mask_path, 'cuda', usr_defined_modules=['SiLU'])
    m_speedup.speedup_model()

    # eport prune model
    model.head.decode_in_inference = False
    traced_model_path = osp.join(cfg.work_dir, model_name + '_traced_model.pt')
    traced_model = torch.jit.trace(model, dummy_input)
    torch.jit.save(traced_model, traced_model_path)

    if cfg.oss_work_dir is not None:
        export_oss_path = os.path.join(cfg.oss_work_dir, 'prune_model.pt')
        if not os.path.exists(traced_model_path):
            logger.warning(f'{traced_model_path} does not exists, skip upload')
        else:
            logger.info(f'upload {traced_model_path} to {export_oss_path}')
            io.safe_copy(traced_model_path, export_oss_path)