-
Notifications
You must be signed in to change notification settings - Fork 205
/
run_cot.py
127 lines (110 loc) · 3.01 KB
/
run_cot.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
123
124
125
126
127
# -*- coding: utf-8 -*-
# @Author : William
# @Project : TextGAN-william
# @FileName : run_cot.py
# @Time : Created at 2020/4/21
# @Blog : http://zhiweil.ml/
# @Description :
# Copyrights (C) 2018. All Rights Reserved.
import sys
from subprocess import call
import os
# Job id and gpu_id
if len(sys.argv) > 2:
job_id = int(sys.argv[1])
gpu_id = str(sys.argv[2])
print('job_id: {}, gpu_id: {}'.format(job_id, gpu_id))
elif len(sys.argv) > 1:
job_id = int(sys.argv[1])
gpu_id = 0
print('job_id: {}, missing gpu_id (use default {})'.format(job_id, gpu_id))
else:
job_id = 0
gpu_id = 0
print('Missing argument: job_id and gpu_id. Use default job_id: {}, gpu_id: {}'.format(job_id, gpu_id))
# Executables
executable = 'python' # specify your own python interpreter path here
rootdir = '../'
scriptname = 'main.py'
# ===Program===
if_test = int(False)
run_model = 'cot'
CUDA = int(True)
oracle_pretrain = int(True)
gen_pretrain = int(False)
dis_pretrain = int(False)
MLE_train_epoch = 0
ADV_train_epoch = 20000
tips = 'CoT experiments'
# ===Oracle or Real===
if_real_data = [int(False), int(True), int(True)]
dataset = ['oracle', 'image_coco', 'emnlp_news']
vocab_size = [5000, 0, 0]
# ===Basic Param===
data_shuffle = int(False)
model_type = 'vanilla'
gen_init = 'normal'
dis_init = 'normal'
samples_num = 10000
batch_size = 64
max_seq_len = 20
gen_lr = 1e-2
pre_log_step = 10
adv_log_step = 100
# ===Generator===
ADV_g_step = 1
gen_embed_dim = 32
gen_hidden_dim = 32
# ===Discriminator===
ADV_d_step = 1
# ===Metrics===
use_nll_oracle = int(True)
use_nll_gen = int(True)
use_nll_div = int(True)
use_bleu = int(True)
use_self_bleu = int(True)
use_ppl = int(False)
args = [
# Program
'--if_test', if_test,
'--run_model', run_model,
'--cuda', CUDA,
# '--device', gpu_id, # comment for auto GPU
'--ora_pretrain', oracle_pretrain,
'--gen_pretrain', gen_pretrain,
'--dis_pretrain', dis_pretrain,
'--mle_epoch', MLE_train_epoch,
'--adv_epoch', ADV_train_epoch,
'--tips', tips,
# Oracle or Real
'--if_real_data', if_real_data[job_id],
'--dataset', dataset[job_id],
'--vocab_size', vocab_size[job_id],
# Basic Param
'--shuffle', data_shuffle,
'--model_type', model_type,
'--gen_init', gen_init,
'--dis_init', dis_init,
'--samples_num', samples_num,
'--batch_size', batch_size,
'--max_seq_len', max_seq_len,
'--gen_lr', gen_lr,
'--pre_log_step', pre_log_step,
'--adv_log_step', adv_log_step,
# Generator
'--adv_g_step', ADV_g_step,
'--gen_embed_dim', gen_embed_dim,
'--gen_hidden_dim', gen_hidden_dim,
# Discriminator
'--adv_d_step', ADV_d_step,
# Metrics
'--use_nll_oracle', use_nll_oracle,
'--use_nll_gen', use_nll_gen,
'--use_nll_div', use_nll_div,
'--use_bleu', use_bleu,
'--use_self_bleu', use_self_bleu,
'--use_ppl', use_ppl,
]
args = list(map(str, args))
my_env = os.environ.copy()
call([executable, scriptname] + args, env=my_env, cwd=rootdir)