Пример #1
0
def main(config):

    # init loaders and base
    loaders = ReIDLoaders(config)
    base = Base(config)

    # make directions
    make_dirs(base.output_path)

    # init logger
    logger = Logger(os.path.join(config.output_path, 'log.txt'))
    logger(config)

    assert config.mode in ['train', 'test', 'visualize']
    if config.mode == 'train':  # train mode

        # automatically resume model from the latest one
        if config.auto_resume_training_from_lastest_steps:
            start_train_epoch = base.resume_last_model()

        # main loop
        for current_epoch in range(start_train_epoch,
                                   config.total_train_epochs):
            # save model
            base.save_model(current_epoch)
            # train
            _, results = train_an_epoch(config, base, loaders, current_epoch)
            logger('Time: {};  Epoch: {};  {}'.format(time_now(),
                                                      current_epoch, results))

        # test
        base.save_model(config.total_train_epochs)
        mAP, CMC, pres, recalls, thresholds = test(config, base, loaders)
        logger('Time: {}; Test Dataset: {}, \nmAP: {} \nRank: {}'.format(
            time_now(), config.test_dataset, mAP, CMC))
        plot_prerecall_curve(config, pres, recalls, thresholds, mAP, CMC,
                             'none')

    elif config.mode == 'test':  # test mode
        base.resume_from_model(config.resume_test_model)
        mAP, CMC, pres, recalls, thresholds = test(config, base, loaders)
        logger('Time: {}; Test Dataset: {}, \nmAP: {} \nRank: {}'.format(
            time_now(), config.test_dataset, mAP, CMC))
        logger(
            'Time: {}; Test Dataset: {}, \nprecision: {} \nrecall: {}\nthresholds: {}'
            .format(time_now(), config.test_dataset, mAP, CMC, pres, recalls,
                    thresholds))
        plot_prerecall_curve(config, pres, recalls, thresholds, mAP, CMC,
                             'none')

    elif config.mode == 'visualize':  # visualization mode
        base.resume_from_model(config.resume_visualize_model)
        visualize(config, base, loaders)
def main(config):

    # init loaders and base
    loaders = ReIDLoaders(config)
    base = Base(config)

    # make directions
    make_dirs(base.output_path)

    # init logger
    logger = Logger(os.path.join(config.output_path, 'log.txt'))
    logger(config)

    assert config.mode in ['train', 'test', 'visualize']
    if config.mode == 'train':  # train mode

        # automatically resume model from the latest one
        if config.auto_resume_training_from_lastest_steps:
            print('resume', base.output_path)
            start_train_epoch = base.resume_last_model()
        #start_train_epoch = 0

        # main loop
        for current_epoch in range(start_train_epoch,
                                   config.total_train_epochs + 1):
            # save model
            base.save_model(current_epoch)
            # train
            base.lr_scheduler.step(current_epoch)
            _, results = train_an_epoch(config, base, loaders)
            logger('Time: {};  Epoch: {};  {}'.format(time_now(),
                                                      current_epoch, results))

        # test
        base.save_model(config.total_train_epochs)
        mAP, CMC = test(config, base, loaders)
        logger('Time: {}; Test Dataset: {}, \nmAP: {} \nRank: {}'.format(
            time_now(), config.test_dataset, mAP, CMC))

    elif config.mode == 'test':  # test mode
        base.resume_from_model(config.resume_test_model)
        mAP, CMC = test(config, base, loaders)
        logger('Time: {}; Test Dataset: {}, \nmAP: {} \nRank: {} with len {}'.
               format(time_now(), config.test_dataset, mAP, CMC, len(CMC)))

    elif config.mode == 'visualize':  # visualization mode
        base.resume_from_model(config.resume_visualize_model)
        visualize(config, base, loaders)