Exemplo n.º 1
0
def main():
    # 1. Get input arguments
    args = get_args()

    # 2. Create config instance from args above
    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available()
    if args.config_file:
        cfg.merge_from_file(args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)
    set_random_seed(cfg.train.seed)

    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))

    print('Show configuration\n{}\n'.format(cfg))
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    # 3. Create DataManager Instance
    datamanager = build_datamanager(cfg)

    print('Building model: {}'.format(cfg.model.name))
    model = torchreid.models.build_model(
        name=cfg.model.name,
        num_classes=datamanager.num_train_pids,
        loss=cfg.loss.name,
        pretrained=cfg.model.pretrained,
        use_gpu=cfg.use_gpu)
    num_params, flops = compute_model_complexity(
        model, (1, 3, cfg.data.height, cfg.data.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
        load_pretrained_weights(model, cfg.model.load_weights)

    if cfg.use_gpu:
        model = nn.DataParallel(model).cuda()

    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
    scheduler = torchreid.optim.build_lr_scheduler(optimizer,
                                                   **lr_scheduler_kwargs(cfg))

    if cfg.model.resume and check_isfile(cfg.model.resume):
        cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
                                                       model,
                                                       optimizer=optimizer,
                                                       scheduler=scheduler)

    print('Building {}-engine for {}-reid'.format(cfg.loss.name,
                                                  cfg.data.type))

    # Build engine and run
    engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
    engine.run(**engine_run_kwargs(cfg))
Exemplo n.º 2
0
def main(args):
    regex_mAP = re.compile(r'mAP: ([\.\deE+-]+)%')
    regex_r1 = re.compile(r'Rank-1  : ([\.\deE+-]+)%')
    regex_r5 = re.compile(r'Rank-5  : ([\.\deE+-]+)%')
    regex_r10 = re.compile(r'Rank-10 : ([\.\deE+-]+)%')
    regex_r20 = re.compile(r'Rank-20 : ([\.\deE+-]+)%')

    final_res = defaultdict(list)

    directories = listdir_nohidden(args.directory, sort=True)
    num_dirs = len(directories)
    for directory in directories:
        fullpath = os.path.join(args.directory, directory)
        filepath = glob.glob(os.path.join(fullpath, 'test.log*'))[0]
        check_isfile(filepath)
        print(f'Parsing {filepath}')
        res = parse_file(filepath, regex_mAP, regex_r1, regex_r5, regex_r10,
                         regex_r20)
        for key, value in res.items():
            final_res[key].append(value)

    print('Finished parsing')
    print(f'The average results over {num_dirs} splits are shown below')

    for key, values in final_res.items():
        mean_val = np.mean(values)
        print(f'{key}: {mean_val:.1f}')
Exemplo n.º 3
0
def build_auxiliary_model(config_file,
                          num_classes,
                          use_gpu,
                          device_ids,
                          num_iter,
                          lr=None,
                          nncf_aux_config_changes=None,
                          aux_config_opts=None,
                          aux_pretrained_dict=None):
    aux_cfg = get_default_config()
    aux_cfg.use_gpu = use_gpu
    merge_from_files_with_base(aux_cfg, config_file)
    if nncf_aux_config_changes:
        print(
            f'applying to aux config changes from NNCF aux config {nncf_aux_config_changes}'
        )
        if not isinstance(nncf_aux_config_changes, CfgNode):
            nncf_aux_config_changes = CfgNode(nncf_aux_config_changes)
        aux_cfg.merge_from_other_cfg(nncf_aux_config_changes)
    if aux_config_opts:
        print(f'applying to aux config changes from command line arguments, '
              f'the changes are:\n{pformat(aux_config_opts)}')
        aux_cfg.merge_from_list(aux_config_opts)

    print(f'\nShow auxiliary configuration\n{aux_cfg}\n')

    if lr is not None:
        aux_cfg.train.lr = lr
        print(f"setting learning rate from main model: {lr}")
    model = torchreid.models.build_model(**model_kwargs(aux_cfg, num_classes))
    optimizer = torchreid.optim.build_optimizer(model,
                                                **optimizer_kwargs(aux_cfg))
    scheduler = torchreid.optim.build_lr_scheduler(
        optimizer=optimizer, num_iter=num_iter, **lr_scheduler_kwargs(aux_cfg))

    if aux_cfg.model.resume and check_isfile(aux_cfg.model.resume):
        aux_cfg.train.start_epoch = resume_from_checkpoint(
            aux_cfg.model.resume,
            model,
            optimizer=optimizer,
            scheduler=scheduler)

    elif aux_pretrained_dict is not None:
        load_pretrained_weights(model, pretrained_dict=aux_pretrained_dict)

    elif aux_cfg.model.load_weights and check_isfile(
            aux_cfg.model.load_weights):
        load_pretrained_weights(model, aux_cfg.model.load_weights)

    if aux_cfg.use_gpu:
        assert device_ids is not None

        if len(device_ids) > 1:
            model = DataParallel(model, device_ids=device_ids,
                                 output_device=0).cuda(device_ids[0])
        else:
            model = model.cuda(device_ids[0])

    return model, optimizer, scheduler
Exemplo n.º 4
0
def main():
    global args

    set_random_seed(args.seed)
    if not args.use_avai_gpus:
        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
    use_gpu = torch.cuda.is_available() and not args.use_cpu
    log_name = 'test.log' if args.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(args.save_dir, log_name))
    print('** Arguments **')
    arg_keys = list(args.__dict__.keys())
    arg_keys.sort()
    for key in arg_keys:
        print('{}: {}'.format(key, args.__dict__[key]))
    print('\n')
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))
    if use_gpu:
        torch.backends.cudnn.benchmark = True
    else:
        warnings.warn(
            'Currently using CPU, however, GPU is highly recommended')

    datamanager = build_datamanager(args)

    print('Building model: {}'.format(args.arch))
    model = torchreid.models.build_model(
        name=args.arch,
        num_classes=datamanager.num_train_pids,
        loss=args.loss.lower(),
        pretrained=(not args.no_pretrained),
        use_gpu=use_gpu)
    num_params, flops = compute_model_complexity(
        model, (1, 3, args.height, args.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if args.load_weights and check_isfile(args.load_weights):
        load_pretrained_weights(model, args.load_weights)

    if use_gpu:
        model = nn.DataParallel(model).cuda()

    optimizer = torchreid.optim.build_optimizer(model,
                                                **optimizer_kwargs(args))

    scheduler = torchreid.optim.build_lr_scheduler(optimizer,
                                                   **lr_scheduler_kwargs(args))

    if args.resume and check_isfile(args.resume):
        args.start_epoch = resume_from_checkpoint(args.resume,
                                                  model,
                                                  optimizer=optimizer)

    print('Building {}-engine for {}-reid'.format(args.loss, args.app))
    engine = build_engine(args, datamanager, model, optimizer, scheduler)

    engine.run(**engine_run_kwargs(args))
Exemplo n.º 5
0
def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config-file', type=str, default='', help='path to config file')
    parser.add_argument('-s', '--sources', type=str, nargs='+', help='source datasets (delimited by space)')
    parser.add_argument('-t', '--targets', type=str, nargs='+', help='target datasets (delimited by space)')
    parser.add_argument('--root', type=str, default='', help='path to data root')
    parser.add_argument('opts', default=None, nargs=argparse.REMAINDER,
                        help='Modify config options using the command-line')
    args = parser.parse_args()

    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available()
    if args.config_file:
        cfg.merge_from_file(args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)
    set_random_seed(cfg.train.seed)

    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))

    print('Show configuration\n{}\n'.format(cfg))
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    datamanager = build_datamanager(cfg)

    print('Building model: {}'.format(cfg.model.name))
    model = torchreid.models.build_model(**model_kwargs(cfg, datamanager.num_train_pids))
    num_params, flops = compute_model_complexity(model, (1, 3, cfg.data.height, cfg.data.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
        if cfg.model.pretrained and not cfg.test.evaluate:
            state_dict = torch.load(cfg.model.load_weights)
            model.load_pretrained_weights(state_dict)
        else:
            load_pretrained_weights(model, cfg.model.load_weights)

    if cfg.use_gpu:
        model = nn.DataParallel(model).cuda()

    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
    scheduler = torchreid.optim.build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg))

    if cfg.model.resume and check_isfile(cfg.model.resume):
        cfg.train.start_epoch = resume_from_checkpoint(
            cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler
        )

    print('Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type))
    engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
    engine.run(**engine_run_kwargs(cfg))
Exemplo n.º 6
0
def main():
    global args

    set_random_seed(args.seed)
    if not args.use_avai_gpus:
        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
    use_gpu = (torch.cuda.is_available() and not args.use_cpu)
    log_name = 'test.log' if args.evaluate else 'train.log'
    sys.stdout = Logger(osp.join(args.save_dir, log_name))
    print('==========\nArgs:{}\n=========='.format(args))
    if use_gpu:
        print('Currently using GPU {}'.format(args.gpu_devices))
        torch.backends.cudnn.benchmark = True
    else:
        warnings.warn(
            'Currently using CPU, however, GPU is highly recommended')

    datamanager = build_datamanager(args)

    print('Building model: {}'.format(args.arch))
    model = torchreid.models.build_model(
        name=args.arch,
        num_classes=datamanager.num_train_pids,
        loss=args.loss.lower(),
        pretrained=(not args.no_pretrained),
        use_gpu=use_gpu)
    num_params, flops = compute_model_complexity(
        model, (1, 3, args.height, args.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if args.load_weights and check_isfile(args.load_weights):
        load_pretrained_weights(model, args.load_weights)

    if use_gpu:
        model = nn.DataParallel(model).cuda()

    optimizer = torchreid.optim.build_optimizer(model,
                                                **optimizer_kwargs(args))

    scheduler = torchreid.optim.build_lr_scheduler(optimizer,
                                                   **lr_scheduler_kwargs(args))

    if args.resume and check_isfile(args.resume):
        args.start_epoch = resume_from_checkpoint(args.resume,
                                                  model,
                                                  optimizer=optimizer)

    print('Building {}-engine for {}-reid'.format(args.loss, args.app))
    engine = build_engine(args, datamanager, model, optimizer, scheduler)

    engine.run(**engine_run_kwargs(args))
Exemplo n.º 7
0
    def __init__(
        self,
        model_name='',
        model_path='',
        image_size=(256, 128),
        pixel_mean=[0.485, 0.456, 0.406],
        pixel_std=[0.229, 0.224, 0.225],
        pixel_norm=True,
        device='cuda',
        verbose=True
    ):
        # Build model
        model = build_model(
            model_name,
            num_classes=1,
            pretrained=not (model_path and check_isfile(model_path)),
            use_gpu=device.startswith('cuda')
        )
        model.eval()

        if verbose:
            num_params, flops = compute_model_complexity(
                model, (1, 3, image_size[0], image_size[1])
            )
            print('Model: {}'.format(model_name))
            print('- params: {:,}'.format(num_params))
            print('- flops: {:,}'.format(flops))

        if model_path and check_isfile(model_path):
            load_pretrained_weights(model, model_path)

        # Build transform functions
        transforms = []
        transforms += [T.Resize(image_size)]
        transforms += [T.ToTensor()]
        if pixel_norm:
            transforms += [T.Normalize(mean=pixel_mean, std=pixel_std)]
        preprocess = T.Compose(transforms)

        to_pil = T.ToPILImage()

        device = torch.device(device)
        model.to(device)

        # Class attributes
        self.model = model
        self.preprocess = preprocess
        self.to_pil = to_pil
        self.device = device
Exemplo n.º 8
0
def main():
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config-file', '-c', type=str, required=True)
    parser.add_argument('--root', '-r', type=str, required=True)
    parser.add_argument('--save-dir', type=str, default='log')
    parser.add_argument('opts', default=None, nargs=REMAINDER)
    args = parser.parse_args()

    assert osp.exists(args.config_file)
    assert osp.exists(args.root)

    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available()
    if args.config_file:
        cfg.merge_from_file(args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    data_loader, num_pids = prepare_data(cfg, mode='gallery')

    print('Building model: {}'.format(cfg.model.name))
    model = torchreid.models.build_model(**model_kwargs(cfg, num_pids))

    if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
        load_pretrained_weights(model, cfg.model.load_weights)

    if cfg.use_gpu:
        model = model.cuda()

    visualize_activation_map(model, data_loader, args.save_dir, cfg.data.width,
                             cfg.data.height, cfg.use_gpu)
Exemplo n.º 9
0
    def __init__(self,
                 config_path='',
                 model_path='',
                 device='cuda',
                 verbose=True):
        # Build model
        cfg = get_default_config()
        merge_from_files_with_base(cfg, config_path)
        cfg.use_gpu = device.startswith('cuda')
        model = build_model(**model_kwargs(cfg, 1))
        model.eval()

        image_size = (cfg.data.height, cfg.data.width)
        flops, num_params = get_model_complexity_info(
            model, (3, image_size[0], image_size[1]),
            as_strings=False,
            verbose=False,
            print_per_layer_stat=False)

        if verbose:
            print('Model: {}'.format(cfg.model.name))
            print('- params: {:,}'.format(num_params))
            print('- flops: {:,}'.format(flops))

        if model_path and check_isfile(model_path):
            load_pretrained_weights(model, model_path)

        # Build transform functions
        transforms = []
        transforms += [T.Resize(image_size)]
        transforms += [T.ToTensor()]
        print(cfg.data.norm_mean, cfg.data.norm_std)
        transforms += [
            T.Normalize(mean=cfg.data.norm_mean, std=cfg.data.norm_std)
        ]
        preprocess = T.Compose(transforms)

        to_pil = T.ToPILImage()

        device = torch.device(device)
        model.to(device)

        # Class attributes
        self.model = model
        self.preprocess = preprocess
        self.to_pil = to_pil
        self.device = device
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--root',
        type=str,
        default=
        '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/V20200108/zhejiang_train'
    )
    parser.add_argument('-d', '--dataset', type=str, default='rock_dataset')
    parser.add_argument('-m', '--model', type=str, default='abd_resnet')
    parser.add_argument('--weights', type=str)
    parser.add_argument('--save-dir', type=str, default='log/resnet50_cam')
    parser.add_argument('--height', type=int, default=672)
    parser.add_argument('--width', type=int, default=672)
    args = parser.parse_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = "3"

    use_gpu = torch.cuda.is_available()
    torchreid.data.register_image_dataset(
        'rock_dataset', torchreid.data.datasets.image.rock_dataset.RockDataSet)

    datamanager = torchreid.data.ImageDataManager(
        root=args.root,
        sources=args.dataset,
        height=args.height,
        width=args.width,
        batch_size_train=4,
        batch_size_test=4,
        transforms=None,
        train_sampler='SequentialSampler')
    test_loader = datamanager.test_loader

    model = torchreid.models.build_model(
        name=args.model,
        num_classes=datamanager.num_train_pids,
        use_gpu=use_gpu)

    if use_gpu:
        model = model.cuda()

    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)

    visactmap(model, test_loader, args.save_dir, args.width, args.height,
              use_gpu)
Exemplo n.º 11
0
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default='')
    parser.add_argument('-m', '--model', type=str, default='resnet50')
    parser.add_argument(
        '--weights',
        type=str,
        default=
        '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/remote_data/PycharmProjects/ABD-Net-master/model_best.pth.tar'
    )
    parser.add_argument('--save-dir', type=str, default='logs/resnet50')
    parser.add_argument('--height', type=int, default=672)
    parser.add_argument('--width', type=int, default=672)
    args = parser.parse_args()

    use_gpu = torch.cuda.is_available()

    test_dir = '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/测试集/ceshi/crop/浙江省温州苍南县西古庵早白垩世小平田组PM201(挑选3张泛化测试用)20200114'
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    test_dataset = datasets.ImageFolder(
        test_dir, transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ]))
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=8,
                                              shuffle=False,
                                              num_workers=8,
                                              pin_memory=True)

    model = torchreid.models.build_model(name=args.model,
                                         num_classes=70,
                                         use_gpu=use_gpu)

    if use_gpu:
        model = model.cuda()

    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)

    visactmap(model, test_loader, args.save_dir, args.width, args.height,
              use_gpu)
Exemplo n.º 12
0
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str)
    parser.add_argument('-d', '--dataset', type=str, default='market1501')
    parser.add_argument('-m', '--model', type=str, default='osnet_x1_0')
    parser.add_argument('--weights', type=str)
    parser.add_argument('--save-dir', type=str, default='log')
    parser.add_argument('--height', type=int, default=256)
    parser.add_argument('--width', type=int, default=128)
    args = parser.parse_args()

    use_gpu = torch.cuda.is_available()

    datamanager = torchreid.data.ImageDataManager(
        root=args.root,
        sources=args.dataset,
        height=args.height,
        width=args.width,
        batch_size_train=100,
        batch_size_test=100,
        transforms=None,
        train_sampler='SequentialSampler'
    )
    test_loader = datamanager.test_loader

    model = torchreid.models.build_model(
        name=args.model,
        num_classes=datamanager.num_train_pids,
        use_gpu=use_gpu
    )

    if use_gpu:
        model = model.cuda()

    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)

    visactmap(
        model, test_loader, args.save_dir, args.width, args.height, use_gpu
    )
Exemplo n.º 13
0
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--root',
        type=str,
        default=
        '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/remote_data/Market1501'
    )
    parser.add_argument('-d', '--dataset', type=str, default='market1501')
    parser.add_argument('-m', '--model', type=str, default='abd_resnet50')
    # /home/ddj2/PycharmProjects/deep-person-reid-master/saves/checkpoint_best.pth.tar
    parser.add_argument(
        '--weights',
        type=str,
        default='/media/ddj2/8b1bfd93-3f3f-4475-b279-6a9ae59c6639/'
        'remote_dir/checkpoint/market_checkpoint_best.pth.tar')
    parser.add_argument(
        '--save-dir',
        type=str,
        default=
        '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/remote_data/Market1501/log/abd'
        '-official-pretrained')
    parser.add_argument('--height', type=int, default=384)
    parser.add_argument('--width', type=int, default=128)

    args = parser.parse_args()

    new_args = {
        'shallow_cam': True,
        'compatibility': False,
        'branches': ['global', 'abd'],
        'abd_dim': 1024,
        'global_dim': 1024,
        'abd_np': 2,
        'abd_dan': ['cam', 'pam'],
        'abd_dan_no_head': False,
        'dropout': 0.5,
        'global_max_pooling': False,
        'use_ow': True,
        'margin': 1.2,
        'label_smooth': True,
        'flip_eval': True,
    }

    use_gpu = torch.cuda.is_available()

    datamanager = torchreid.data.ImageDataManager(
        root=args.root,
        sources=args.dataset,
        height=args.height,
        width=args.width,
        batch_size_train=100,
        batch_size_test=100,
        transforms=None,
        train_sampler='SequentialSampler')
    test_loader = datamanager.test_loader

    # model = resnet_orig.resnet50(num_classes=1501)
    model = torchreid.models.build_model(
        name=args.model,
        num_classes=datamanager.num_train_pids,
        use_gpu=use_gpu,
        args=new_args)

    if use_gpu:
        model = model.cuda()

    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)

    visactmap(model, test_loader, args.save_dir, args.width, args.height,
              use_gpu)
Exemplo n.º 14
0
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        '--config-file', type=str, default='', help='path to config file'
    )
    parser.add_argument(
        '-s',
        '--sources',
        type=str,
        nargs='+',
        help='source datasets (delimited by space)'
    )
    parser.add_argument(
        '-t',
        '--targets',
        type=str,
        nargs='+',
        help='target datasets (delimited by space)'
    )
    parser.add_argument(
        '--transforms', type=str, nargs='+', help='data augmentation'
    )
    parser.add_argument(
        '--root', type=str, default='', help='path to data root'
    )
    parser.add_argument(
        '--gpu-devices',
        type=str,
        default='',
    )
    parser.add_argument(
        'opts',
        default=None,
        nargs=argparse.REMAINDER,
        help='Modify config options using the command-line'
    )
    args = parser.parse_args()

    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available()
    if args.config_file:
        cfg.merge_from_file(args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)
    set_random_seed(cfg.train.seed)

    if cfg.use_gpu and args.gpu_devices:
        # if gpu_devices is not specified, all available gpus will be used
        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))

    print('Show configuration\n{}\n'.format(cfg))
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))

    print('Building model: {}'.format(cfg.model.name))
    model = osnet_models.build_model(
        cfg.model.name, num_classes=datamanager.num_train_pids
    )
    num_params, flops = compute_model_complexity(
        model, (1, 3, cfg.data.height, cfg.data.width)
    )
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if cfg.use_gpu:
        model = nn.DataParallel(model).cuda()

    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
    scheduler = torchreid.optim.build_lr_scheduler(
        optimizer, **lr_scheduler_kwargs(cfg)
    )

    if cfg.model.resume and check_isfile(cfg.model.resume):
        cfg.train.start_epoch = resume_from_checkpoint(
            cfg.model.resume, model, optimizer=optimizer
        )

    print('Building NAS engine')
    engine = ImageSoftmaxNASEngine(
        datamanager,
        model,
        optimizer,
        scheduler=scheduler,
        use_gpu=cfg.use_gpu,
        label_smooth=cfg.loss.softmax.label_smooth,
        mc_iter=cfg.nas.mc_iter,
        init_lmda=cfg.nas.init_lmda,
        min_lmda=cfg.nas.min_lmda,
        lmda_decay_step=cfg.nas.lmda_decay_step,
        lmda_decay_rate=cfg.nas.lmda_decay_rate,
        fixed_lmda=cfg.nas.fixed_lmda
    )
    engine.run(**engine_run_kwargs(cfg))

    print('*** Display the found architecture ***')
    if cfg.use_gpu:
        model.module.build_child_graph()
    else:
        model.build_child_graph()
Exemplo n.º 15
0
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--model-name',
                        type=str,
                        default='',
                        help='Model name')
    parser.add_argument('--weights', type=str, default='', help='Weights path')
    parser.add_argument('--output',
                        type=str,
                        default='output',
                        help='Output path')
    parser.add_argument('--resolution',
                        type=str,
                        default='128x256',
                        help='Resolution (WxH)')

    args = parser.parse_args()
    width, height = [int(i) for i in args.resolution.split('x')]

    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    imagedata_kwargs = {
        'root': 'reid-data',
        'sources': ['market1501'],
        'targets': ['market1501'],
        'height': 256,
        'width': 128,
        'transforms': ['random_flip', 'color_jitter'],
        'norm_mean': [0.485, 0.456, 0.406],
        'norm_std': [0.229, 0.224, 0.225],
        'use_gpu': False,
        'split_id': 0,
        'combineall': False,
        'load_train_targets': False,
        'batch_size_train': 64,
        'batch_size_test': 300,
        'workers': 4,
        'num_instances': 4,
        'train_sampler': 'RandomSampler',
        'cuhk03_labeled': False,
        'cuhk03_classic_split': False,
        'market1501_500k': False
    }
    datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs)

    print('Building model: {}'.format(args.model_name))
    model = torchreid.models.build_model(
        name=args.model_name,
        num_classes=datamanager.num_train_pids,
        loss='softmax',
        pretrained=True,
        use_gpu=False)
    num_params, flops = compute_model_complexity(model, (1, 3, height, width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)
    _input = torch.Tensor(1, 3, height, width)
    inputs = (_input, )

    print('Converting PyTorch model to ONNX...')
    tmp = tempfile.mktemp(suffix='.onnx')
    torch.onnx._export(model, inputs, tmp, export_params=True)

    onnx_model = onnx.load(tmp)
    export_path = args.output

    onnx.checker.check_model(onnx_model)

    print('Prepare TF model...')
    tf_rep = prepare(onnx_model, strict=False)

    if path.exists(export_path):
        shutil.rmtree(export_path)

    with tf.Session() as persisted_sess:
        print("load graph")
        persisted_sess.graph.as_default()
        tf.import_graph_def(tf_rep.graph.as_graph_def(), name='')

        i_tensors = []
        o_tensors = []
        inputs = {}
        outputs = {}

        for i in tf_rep.inputs:
            t = persisted_sess.graph.get_tensor_by_name(
                tf_rep.tensor_dict[i].name)
            i_tensors.append(t)
            tensor_info = tf.saved_model.utils.build_tensor_info(t)
            inputs[t.name.split(':')[0].lower()] = tensor_info
            print('input tensor [name=%s, type=%s, shape=%s]' %
                  (t.name, t.dtype.name, t.shape.as_list()))
        print('')

        for i in tf_rep.outputs:
            t = persisted_sess.graph.get_tensor_by_name(
                tf_rep.tensor_dict[i].name)
            o_tensors.append(t)
            tensor_info = tf.saved_model.utils.build_tensor_info(t)
            outputs[t.name.split(':')[0]] = tensor_info
            print('output tensor [name=%s, type=%s, shape=%s]' %
                  (t.name, t.dtype.name, t.shape.as_list()))

        feed_dict = {}
        for i in i_tensors:
            feed_dict[i] = np.random.rand(*i.shape.as_list()).astype(
                i.dtype.name)

        print('test run:')
        res = persisted_sess.run(o_tensors, feed_dict=feed_dict)
        print(res)

        # print('INPUTS')
        # print(inputs)
        # print('OUTPUTS')
        # print(outputs)
        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs=inputs,
                outputs=outputs,
                method_name=tf.saved_model.signature_constants.
                PREDICT_METHOD_NAME))
        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
        builder.add_meta_graph_and_variables(
            persisted_sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                prediction_signature
            })
        builder.save()
        print('Model saved to %s' % export_path)
Exemplo n.º 16
0
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config-file',
                        type=str,
                        default='',
                        help='path to config file')
    parser.add_argument('-s',
                        '--sources',
                        type=str,
                        nargs='+',
                        help='source datasets (delimited by space)')
    parser.add_argument('-t',
                        '--targets',
                        type=str,
                        nargs='+',
                        help='target datasets (delimited by space)')
    parser.add_argument('--transforms',
                        type=str,
                        nargs='+',
                        help='data augmentation')
    parser.add_argument('--root',
                        type=str,
                        default='',
                        help='path to data root')
    parser.add_argument('opts',
                        default=None,
                        nargs=argparse.REMAINDER,
                        help='Modify config options using the command-line')
    args = parser.parse_args()

    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available()
    if args.config_file:
        cfg.merge_from_file(args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)
    set_random_seed(cfg.train.seed)

    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))

    print('Show configuration\n{}\n'.format(cfg))
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))

    print('Building model-1: {}'.format(cfg.model.name))
    model1 = torchreid.models.build_model(
        name=cfg.model.name,
        num_classes=datamanager.num_train_pids,
        loss=cfg.loss.name,
        pretrained=cfg.model.pretrained,
        use_gpu=cfg.use_gpu)
    num_params, flops = compute_model_complexity(
        model1, (1, 3, cfg.data.height, cfg.data.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    print('Copying model-1 to model-2')
    model2 = copy.deepcopy(model1)

    if cfg.model.load_weights1 and check_isfile(cfg.model.load_weights1):
        load_pretrained_weights(model1, cfg.model.load_weights1)

    if cfg.model.load_weights2 and check_isfile(cfg.model.load_weights2):
        load_pretrained_weights(model2, cfg.model.load_weights2)

    if cfg.use_gpu:
        model1 = nn.DataParallel(model1).cuda()
        model2 = nn.DataParallel(model2).cuda()

    optimizer1 = torchreid.optim.build_optimizer(model1,
                                                 **optimizer_kwargs(cfg))
    scheduler1 = torchreid.optim.build_lr_scheduler(optimizer1,
                                                    **lr_scheduler_kwargs(cfg))

    optimizer2 = torchreid.optim.build_optimizer(model2,
                                                 **optimizer_kwargs(cfg))
    scheduler2 = torchreid.optim.build_lr_scheduler(optimizer2,
                                                    **lr_scheduler_kwargs(cfg))

    if cfg.model.resume1 and check_isfile(cfg.model.resume1):
        cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume1,
                                                       model1,
                                                       optimizer=optimizer1,
                                                       scheduler=scheduler1)

    if cfg.model.resume2 and check_isfile(cfg.model.resume2):
        resume_from_checkpoint(cfg.model.resume2,
                               model2,
                               optimizer=optimizer2,
                               scheduler=scheduler2)

    print('Building DML-engine for image-reid')
    engine = ImageDMLEngine(datamanager,
                            model1,
                            optimizer1,
                            scheduler1,
                            model2,
                            optimizer2,
                            scheduler2,
                            margin=cfg.loss.triplet.margin,
                            weight_t=cfg.loss.triplet.weight_t,
                            weight_x=cfg.loss.triplet.weight_x,
                            weight_ml=cfg.loss.dml.weight_ml,
                            use_gpu=cfg.use_gpu,
                            label_smooth=cfg.loss.softmax.label_smooth,
                            deploy=cfg.model.deploy)
    engine.run(**engine_run_kwargs(cfg))
Exemplo n.º 17
0
def main():
    parser = build_base_argparser()
    parser.add_argument('-e',
                        '--auxiliary-models-cfg',
                        type=str,
                        nargs='*',
                        default='',
                        help='path to extra config files')
    parser.add_argument('--split-models',
                        action='store_true',
                        help='whether to split models on own gpu')
    parser.add_argument('--enable_quantization',
                        action='store_true',
                        help='Enable NNCF quantization algorithm')
    parser.add_argument('--enable_pruning',
                        action='store_true',
                        help='Enable NNCF pruning algorithm')
    parser.add_argument(
        '--aux-config-opts',
        nargs='+',
        default=None,
        help='Modify aux config options using the command-line')
    args = parser.parse_args()

    cfg = get_default_config()
    cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0
    if args.config_file:
        merge_from_files_with_base(cfg, args.config_file)
    reset_config(cfg, args)
    cfg.merge_from_list(args.opts)

    is_nncf_used = args.enable_quantization or args.enable_pruning
    if is_nncf_used:
        print(f'Using NNCF -- making NNCF changes in config')
        cfg = make_nncf_changes_in_config(cfg, args.enable_quantization,
                                          args.enable_pruning, args.opts)

    set_random_seed(cfg.train.seed, cfg.train.deterministic)

    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))

    print('Show configuration\n{}\n'.format(cfg))

    if cfg.use_gpu:
        torch.backends.cudnn.benchmark = True

    num_aux_models = len(cfg.mutual_learning.aux_configs)
    datamanager = build_datamanager(cfg, args.classes)
    num_train_classes = datamanager.num_train_pids

    print('Building main model: {}'.format(cfg.model.name))
    model = torchreid.models.build_model(
        **model_kwargs(cfg, num_train_classes))
    macs, num_params = get_model_complexity_info(
        model, (3, cfg.data.height, cfg.data.width),
        as_strings=False,
        verbose=False,
        print_per_layer_stat=False)
    print('Main model complexity: params={:,} flops={:,}'.format(
        num_params, macs * 2))

    aux_lr = cfg.train.lr  # placeholder, needed for aux models, may be filled by nncf part below
    if is_nncf_used:
        print('Begin making NNCF changes in model')
        if cfg.use_gpu:
            model.cuda()

        compression_ctrl, model, cfg, aux_lr, nncf_metainfo = \
            make_nncf_changes_in_training(model, cfg,
                                          args.classes,
                                          args.opts)

        should_freeze_aux_models = True
        print(f'should_freeze_aux_models = {should_freeze_aux_models}')
        print('End making NNCF changes in model')
    else:
        compression_ctrl = None
        should_freeze_aux_models = False
        nncf_metainfo = None
    # creating optimizer and scheduler -- it should be done after NNCF part, since
    # NNCF could change some parameters
    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))

    if cfg.lr_finder.enable and not cfg.model.resume:
        scheduler = None
    else:
        scheduler = torchreid.optim.build_lr_scheduler(
            optimizer=optimizer,
            num_iter=datamanager.num_iter,
            **lr_scheduler_kwargs(cfg))
    # Loading model (and optimizer and scheduler in case of resuming training).
    # Note that if NNCF is used, loading is done inside NNCF part, so loading here is not required.
    if cfg.model.resume and check_isfile(
            cfg.model.resume) and not is_nncf_used:
        device_ = 'cuda' if cfg.use_gpu else 'cpu'
        cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
                                                       model,
                                                       optimizer=optimizer,
                                                       scheduler=scheduler,
                                                       device=device_)
    elif cfg.model.load_weights and not is_nncf_used:
        load_pretrained_weights(model, cfg.model.load_weights)

    if cfg.model.type == 'classification':
        check_classification_classes(model,
                                     datamanager,
                                     args.classes,
                                     test_only=cfg.test.evaluate)

    model, extra_device_ids = put_main_model_on_the_device(
        model, cfg.use_gpu, args.gpu_num, num_aux_models, args.split_models)

    if cfg.lr_finder.enable and not cfg.test.evaluate and not cfg.model.resume:
        aux_lr, model, optimizer, scheduler = run_lr_finder(
            cfg,
            datamanager,
            model,
            optimizer,
            scheduler,
            args.classes,
            rebuild_model=True,
            gpu_num=args.gpu_num,
            split_models=args.split_models)

    log_dir = cfg.data.tb_log_dir if cfg.data.tb_log_dir else cfg.data.save_dir
    run_training(cfg,
                 datamanager,
                 model,
                 optimizer,
                 scheduler,
                 extra_device_ids,
                 aux_lr,
                 tb_writer=SummaryWriter(log_dir=log_dir),
                 should_freeze_aux_models=should_freeze_aux_models,
                 nncf_metainfo=nncf_metainfo,
                 compression_ctrl=compression_ctrl)
Exemplo n.º 18
0
def main():
    global args

    set_random_seed(args.seed)
    use_gpu = torch.cuda.is_available() and not args.use_cpu
    log_name = 'test.log' if args.evaluate else 'train.log'
    sys.stdout = Logger(osp.join(args.save_dir, log_name))

    print('** Arguments **')
    arg_keys = list(args.__dict__.keys())
    arg_keys.sort()
    for key in arg_keys:
        print('{}: {}'.format(key, args.__dict__[key]))
    print('\n')
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))

    if use_gpu:
        torch.backends.cudnn.benchmark = True
    else:
        warnings.warn(
            'Currently using CPU, however, GPU is highly recommended')

    dataset_vars = init_dataset(use_gpu)
    trainloader, valloader, testloader, num_attrs, attr_dict = dataset_vars

    if args.weighted_bce:
        print('Use weighted binary cross entropy')
        print('Computing the weights ...')
        bce_weights = torch.zeros(num_attrs, dtype=torch.float)
        for _, attrs, _ in trainloader:
            bce_weights += attrs.sum(0)  # sum along the batch dim
        bce_weights /= len(trainloader) * args.batch_size
        print('Sample ratio for each attribute: {}'.format(bce_weights))
        bce_weights = torch.exp(-1 * bce_weights)
        print('BCE weights: {}'.format(bce_weights))
        bce_weights = bce_weights.expand(args.batch_size, num_attrs)
        criterion = nn.BCEWithLogitsLoss(weight=bce_weights)

    else:
        print('Use plain binary cross entropy')
        criterion = nn.BCEWithLogitsLoss()

    print('Building model: {}'.format(args.arch))
    model = models.build_model(args.arch,
                               num_attrs,
                               pretrained=not args.no_pretrained,
                               use_gpu=use_gpu)
    num_params, flops = compute_model_complexity(
        model, (1, 3, args.height, args.width))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if args.load_weights and check_isfile(args.load_weights):
        load_pretrained_weights(model, args.load_weights)

    if use_gpu:
        model = nn.DataParallel(model).cuda()
        criterion = criterion.cuda()

    if args.evaluate:
        test(model, testloader, attr_dict, use_gpu)
        return

    optimizer = torchreid.optim.build_optimizer(model,
                                                **optimizer_kwargs(args))
    scheduler = torchreid.optim.build_lr_scheduler(optimizer,
                                                   **lr_scheduler_kwargs(args))

    start_epoch = args.start_epoch
    best_result = -np.inf
    if args.resume and check_isfile(args.resume):
        checkpoint = torch.load(args.resume)
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        best_result = checkpoint['label_mA']
        print('Loaded checkpoint from "{}"'.format(args.resume))
        print('- start epoch: {}'.format(start_epoch))
        print('- label_mA: {}'.format(best_result))

    time_start = time.time()

    for epoch in range(start_epoch, args.max_epoch):
        train(epoch, model, criterion, optimizer, scheduler, trainloader,
              use_gpu)
        test_outputs = test(model, testloader, attr_dict, use_gpu)
        label_mA = test_outputs[0]
        is_best = label_mA > best_result
        if is_best:
            best_result = label_mA

        save_checkpoint(
            {
                'state_dict': model.state_dict(),
                'epoch': epoch + 1,
                'label_mA': label_mA,
                'optimizer': optimizer.state_dict(),
            },
            args.save_dir,
            is_best=is_best)

    elapsed = round(time.time() - time_start)
    elapsed = str(datetime.timedelta(seconds=elapsed))
    print('Elapsed {}'.format(elapsed))
Exemplo n.º 19
0
def objective(cfg, args, trial):
    # Generate the trials.
    # g_ = trial.suggest_int("g_", 1, 7)
    # asl_pm = trial.suggest_float("asl_pm", 0, 0.5)
    # m = trial.suggest_float("m", 0.01, 0.7)
    # s = trial.suggest_int("s", 5, 60)
    lr = trial.suggest_float("lr", 0.001, 0.5)
    # t = trial.suggest_int("t", 1, 7)
    # cfg.loss.softmax.m = m
    # cfg.loss.softmax.s = s
    # cfg.loss.asl.p_m = asl_pm
    # cfg.loss.am_binary.amb_t = t
    cfg.train.lr = lr

    # geterate damanager
    num_aux_models = len(cfg.mutual_learning.aux_configs)
    datamanager = build_datamanager(cfg, args.classes)

    # build the model.
    num_train_classes = datamanager.num_train_pids
    print('Building main model: {}'.format(cfg.model.name))
    model = torchreid.models.build_model(
        **model_kwargs(cfg, num_train_classes))
    aux_lr = cfg.train.lr  # placeholder, needed for aux models, may be filled by nncf part below
    compression_ctrl = None
    should_freeze_aux_models = False
    nncf_metainfo = None
    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
    scheduler = torchreid.optim.build_lr_scheduler(
        optimizer=optimizer,
        num_iter=datamanager.num_iter,
        **lr_scheduler_kwargs(cfg))
    # Loading model (and optimizer and scheduler in case of resuming training).
    if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
        load_pretrained_weights(model, cfg.model.load_weights)

    if cfg.model.type == 'classification':
        check_classification_classes(model,
                                     datamanager,
                                     args.classes,
                                     test_only=cfg.test.evaluate)

    model, extra_device_ids = put_main_model_on_the_device(
        model, cfg.use_gpu, args.gpu_num, num_aux_models, args.split_models)

    num_aux_models = len(cfg.mutual_learning.aux_configs)
    num_train_classes = datamanager.num_train_pids

    if num_aux_models > 0:
        print(
            f'Enabled mutual learning between {len(cfg.mutual_learning.aux_configs) + 1} models.'
        )

        models, optimizers, schedulers = [model], [optimizer], [scheduler]
        for config_file, device_ids in zip(cfg.mutual_learning.aux_configs,
                                           extra_device_ids):
            aux_model, aux_optimizer, aux_scheduler = build_auxiliary_model(
                config_file,
                num_train_classes,
                cfg.use_gpu,
                device_ids,
                num_iter=datamanager.num_iter,
                lr=aux_lr,
                aux_config_opts=args.aux_config_opts)

            models.append(aux_model)
            optimizers.append(aux_optimizer)
            schedulers.append(aux_scheduler)
    else:
        models, optimizers, schedulers = model, optimizer, scheduler
    print(f'Building {cfg.loss.name}-engine')
    engine = build_engine(cfg,
                          datamanager,
                          models,
                          optimizers,
                          schedulers,
                          should_freeze_aux_models=should_freeze_aux_models,
                          nncf_metainfo=nncf_metainfo,
                          compression_ctrl=compression_ctrl,
                          initial_lr=aux_lr)
    test_acc = AverageMeter()
    obj = 0
    engine.start_epoch = 0
    engine.max_epoch = args.epochs
    print(f"\nnext trial with [lr: {lr}]")

    for engine.epoch in range(args.epochs):
        np.random.seed(cfg.train.seed + engine.epoch)
        avg_loss = engine.train(print_freq=20000,
                                fixbase_epoch=0,
                                open_layers=None,
                                lr_finder=False,
                                perf_monitor=None,
                                stop_callback=None)

        top1, _ = engine.test(
            engine.epoch,
            lr_finder=False,
        )

        test_acc.update(top1)
        smooth_top1 = test_acc.avg
        target_metric = smooth_top1 if engine.target_metric == 'test_acc' else avg_loss

        obj = top1
        if not engine.per_batch_annealing:
            engine.update_lr(output_avg_metric=target_metric)

        trial.report(obj, engine.epoch)

        # Handle pruning based on the intermediate value.
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()

        should_exit, _ = engine.exit_on_plateau_and_choose_best(
            top1, smooth_top1)
        should_exit = engine.early_stoping and should_exit
        if should_exit:
            break

    return obj
Exemplo n.º 20
0
def main():

    # Load model configuration
    parser = argparse.ArgumentParser()
    parser.add_argument('-c',
                        '--config',
                        required=True,
                        help='path to configuration file')
    args = parser.parse_args()
    with open(args.config, "r") as ymlfile:
        config = yaml.load(ymlfile, Loader=yaml.FullLoader)

    # Automatically add sub-folder name to config["save_dir"], with the same name
    # as the config file. For example, config["save_dir"] is typically "logs",
    # so this would change config["save_dir"] to "logs/exp01", for example, so that
    # we don't need to change the save_dir in every single config file (it instead
    # automatically generates it from the name of the config file).
    experiment_number = pathlib.Path(args.config).stem
    config["save_dir"] = os.path.join(config["save_dir"], experiment_number)

    # Set random seeds
    set_random_seed(config["seed"])

    # Set up GPU
    if not config["use_avai_gpus"]:
        os.environ['CUDA_VISIBLE_DEVICES'] = config["gpu_devices"]
    use_gpu = torch.cuda.is_available() and not config["use_cpu"]

    # Set up log files
    log_name = 'test.log' if config["evaluate"] else 'train.log'
    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
    sys.stdout = Logger(osp.join(config["save_dir"], log_name))

    # Prepare for training
    print('==========\nArgs:{}\n=========='.format(config))
    print('Collecting env info ...')
    print('** System info **\n{}\n'.format(collect_env_info()))
    if use_gpu:
        torch.backends.cudnn.benchmark = True
    else:
        warnings.warn(
            'Currently using CPU, however, GPU is highly recommended')

    # Build datamanager and model
    datamanager = build_datamanager(config)

    print('Building model: {}'.format(config["arch"]))
    model = torchreid.models.build_model(
        name=config["arch"],
        num_classes=datamanager.num_train_pids,
        loss=config["loss"].lower(),
        pretrained=(not config["no_pretrained"]),
        use_gpu=use_gpu)

    # Compute model complexity
    num_params, flops = compute_model_complexity(
        model, (1, 3, config["height"], config["width"]))
    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    # Load pretrained weights if necessary
    if config["load_weights"] and check_isfile(config["load_weights"]):
        load_pretrained_weights(model, config["load_weights"])

    # Set up multi-gpu
    if use_gpu:
        model = nn.DataParallel(model).cuda()

    # Model settings
    optimizer = torchreid.optim.build_optimizer(model,
                                                **optimizer_kwargs(config))
    scheduler = torchreid.optim.build_lr_scheduler(
        optimizer, **lr_scheduler_kwargs(config))

    if config["resume"] and check_isfile(config["resume"]):
        config["start_epoch"] = resume_from_checkpoint(config["resume"],
                                                       model,
                                                       optimizer=optimizer)

    print('Building {}-engine for {}-reid'.format(config["loss"],
                                                  config["app"]))
    engine = build_engine(config, datamanager, model, optimizer, scheduler)

    engine.run(**engine_run_kwargs(config))