-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
executable file
·122 lines (103 loc) · 6.36 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
# -*- coding:utf-8 -*-
# Created by LuoJie at 12/7/19
from pgn_tf2.eval import evaluate
from pgn_tf2.test import test_and_save, predict_result
from pgn_tf2.train import train
import argparse
from utils.config import train_x_seg_path, train_y_seg_path, test_x_seg_path, sample_total, batch_size, save_result_dir, \
vocab_path, epochs, checkpoint_dir, checkpoint_prefix, vocab_size, val_x_seg_path, val_y_seg_path
from utils.file_utils import get_result_filename
import os
def main():
# 获得参数
parser = argparse.ArgumentParser()
parser.add_argument("--mode", default='eval', help="run mode", type=str)
parser.add_argument("--max_enc_len", default=400, help="Encoder input max sequence length", type=int)
parser.add_argument("--max_dec_len", default=100, help="Decoder input max sequence length", type=int)
parser.add_argument("--batch_size", default=batch_size, help="batch size", type=int)
parser.add_argument("--epochs", default=epochs, help="train epochs", type=int)
parser.add_argument("--vocab_path", default=vocab_path, help="vocab path", type=str)
parser.add_argument("--learning_rate", default=0.15, help="Learning rate", type=float)
parser.add_argument("--adagrad_init_acc", default=0.1,
help="Adagrad optimizer initial accumulator value. "
"Please refer to the Adagrad optimizer API documentation "
"on tensorflow site for more details.",
type=float)
parser.add_argument('--rand_unif_init_mag', default=0.02,
help='magnitude for lstm cells random uniform inititalization', type=float)
parser.add_argument('--eps', default=1e-12, help='eps',
type=float)
parser.add_argument('--trunc_norm_init_std', default=1e-4, help='std of trunc norm init, '
'used for initializing everything else',
type=float)
parser.add_argument('--cov_loss_wt', default=1.0, help='Weight of coverage loss (lambda in the paper).'
' If zero, then no incentive to minimize coverage loss.',
type=float)
parser.add_argument('--max_grad_norm', default=2.0, help='for gradient clipping', type=float)
parser.add_argument("--vocab_size", default=50000, help="max vocab size , None-> Max ", type=int)
parser.add_argument("--max_vocab_size", default=50000, help="max vocab size , None-> Max ", type=int)
parser.add_argument("--beam_size", default=batch_size,
help="beam size for beam search decoding (must be equal to batch size in decode mode)",
type=int)
parser.add_argument("--embed_size", default=300, help="Words embeddings dimension", type=int)
parser.add_argument("--enc_units", default=128, help="Encoder GRU cell units number", type=int)
parser.add_argument("--dec_units", default=256, help="Decoder GRU cell units number", type=int)
parser.add_argument("--attn_units", default=256, help="[context vector, decoder state, decoder input] feedforward \
result dimension - this result is used to compute the attention weights",
type=int)
parser.add_argument("--train_seg_x_dir", default=train_x_seg_path, help="train_seg_x_dir", type=str)
parser.add_argument("--train_seg_y_dir", default=train_y_seg_path, help="train_seg_y_dir", type=str)
parser.add_argument("--val_seg_x_dir", default=val_x_seg_path, help="val_x_seg_path", type=str)
parser.add_argument("--val_seg_y_dir", default=val_y_seg_path, help="val_y_seg_path", type=str)
parser.add_argument("--test_seg_x_dir", default=test_x_seg_path, help="train_seg_x_dir", type=str)
parser.add_argument("--checkpoint_dir", default=checkpoint_dir,
help="checkpoint_dir",
type=str)
parser.add_argument("--checkpoints_save_steps", default=5, help="Save checkpoints every N steps", type=int)
parser.add_argument("--min_dec_steps", default=4, help="min_dec_steps", type=int)
parser.add_argument("--max_train_steps", default=500000 / (batch_size / 8), help="max_train_steps", type=int)
# parser.add_argument("--max_train_steps", default=50, help="max_train_steps", type=int)
parser.add_argument("--save_batch_train_data", default=False, help="save batch train data to pickle", type=bool)
parser.add_argument("--load_batch_train_data", default=False, help="load batch train data from pickle",
type=bool)
parser.add_argument("--test_save_dir", default=save_result_dir, help="test_save_dir", type=str)
parser.add_argument("--pointer_gen", default=True, help="training, eval or test options", type=bool)
parser.add_argument("--use_coverage", default=True, help="test_save_dir", type=bool)
parser.add_argument("--greedy_decode", default=False, help="greedy_decode", type=bool)
parser.add_argument("--result_save_path", default=get_result_filename(batch_size, epochs, 200, 300),
help='result_save_path', type=str)
parser.add_argument("--max_num_to_eval", default=5, help="max_num_to_eval", type=int)
parser.add_argument("--num_to_test", default=20000, help="num_to_test", type=int)
parser.add_argument("--gpu_memory", default=6, help="gpu_memory GB", type=int)
args = parser.parse_args()
params = vars(args)
# print(params)
if params["mode"] == "train":
train(params)
elif params["mode"] == "test":
params['beam_size'] = 2
params['batch_size'] = 2
result_save_path = params['result_save_path']
predict_result(params, result_save_path)
# test_and_save(params)
elif params["mode"] == "eval":
evaluate(params)
elif params['mode'] == 'auto':
# PGN training
params['mode'] = 'train'
# params['use_coverage'] = False
# params['epochs'] = 30
params['use_coverage'] = True
params['epochs'] = 30
train(params)
# predict result
params['mode'] = 'test'
params['beam_size'] = 2
params['batch_size'] = 2
result_save_path = params['result_save_path']
predict_result(params, result_save_path)
# evaluate
params['mode'] = 'eval'
evaluate(params)
if __name__ == '__main__':
main()