main.py 文件源码

def main():
    if len(sys.argv) == 3:
        config = Config(sys.argv[1], sys.argv[2])
    else:
        assert False

    phase = config.items['phase']

    glog.info('generating model...')
    from model_me import Model

    with tf.device('/cpu:0'):
    # with tf.device('/gpu:%d'%config.items['gpu']):
        model = Model()

    # try:
    #     config.items['starting'] = int(config.items['model'].split('_')[-1])
    # except:
    config.items['starting'] = 0

    # snapshot path
    mkdir_safe(config.items['snap_path'])
    mkdir_safe(os.path.join(config.items['snap_path'], 'output_dev'))
    mkdir_safe(os.path.join(config.items['snap_path'], 'output_test'))

    sess_config = tf.ConfigProto(device_count = {'GPU': 0})
    # sess_config = tf.ConfigProto(allow_soft_placement=True)
    # sess_config.gpu_options.allow_growth = True
    from reader import Reader

    train_set = Reader(phase='train', batch_size=config.items['batch_size'], do_shuffle=True, resample=True, distortion=True)
    valid_set = None#Reader(phase='dev', batch_size=1, do_shuffle=False, resample=False, distortion=False)
    test_set = None#Reader(phase='test', batch_size=1, do_shuffle=False, resample=False, distortion=False)

    with tf.Session(config=sess_config) as sess:
        tf.global_variables_initializer().run()
        if 'model' in config.items.keys():
            model.assign_from_pkl(config.items['model'])
            pdb.set_trace()

            glog.info('loading model: %s...' % config.items['model'])
        if phase == 'ctc':
            # model.make_functions()
            glog.info('ctc training...')
            train_valid(sess, model, train_set, valid_set, test_set, config)
        elif phase == 'extract_feature':
            pass
        elif phase == 'get_prediction':
            from reader import Reader
            train_set = Reader(phase='train', batch_size=config.items['batch_size'], do_shuffle=False, resample=False, distortion=False)
            glog.info('feature extracting...')
            get_prediction(model, train_set, config)
        elif phase == 'top_k_prediction':
            from reader import Reader
            train_set = Reader(phase='test', batch_size=config.items['batch_size'], do_shuffle=False, resample=False, distortion=False)
            glog.info('feature extracting...')
            get_top_k_prediction(model, train_set, config)

    glog.info('end')