-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
294 lines (229 loc) · 10.6 KB
/
train.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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
'''
This script handles the training process.
'''
import argparse
import math
import time
import dill as pickle
from tqdm import tqdm
import numpy as np
import random
import os
import torch
import torch.optim as optim
from model import Transformer
from utils import ScheduledOptim
from loss import cal_performance
from data import *
from config import Constants
def patch_src(src, pad_idx):
src = src.permute(0, 1)
return src
def patch_trg(trg, pad_idx):
trg = trg.permute(0, 1)
trg, gold = trg[:, :-1], trg[:, 1:].contiguous().view(-1)
return trg, gold
def train_epoch(model, training_data, optimizer, opt, device, smoothing):
''' Epoch operation in training phase'''
model.train()
total_loss, n_word_total, n_word_correct = 0, 0, 0
desc = ' - (Training) '
for batch in tqdm(training_data, mininterval=2, desc=desc, leave=False):
# prepare data
src_seq = patch_src(batch['src'], opt.src_pad_idx).to(device)
trg_seq, gold = map(lambda x: x.to(device), patch_trg(batch['trg'], opt.trg_pad_idx))
# forward
optimizer.zero_grad()
pred = model(src_seq, trg_seq)
# backward and update parameters
loss, n_correct, n_word = cal_performance(pred, gold, opt.trg_pad_idx, smoothing=smoothing)
loss.backward()
optimizer.step_and_update_lr()
# note keeping
n_word_total += n_word
n_word_correct += n_correct
total_loss += loss.item()
loss_per_word = total_loss / n_word_total
accuracy = n_word_correct / n_word_total
return loss_per_word, accuracy
def eval_epoch(model, validation_data, device, opt):
''' Epoch operation in evaluation phase '''
model.eval()
total_loss, n_word_total, n_word_correct = 0, 0, 0
desc = ' - (Validation) '
with torch.no_grad():
for batch in tqdm(validation_data, mininterval=2, desc=desc, leave=False):
# prepare data
src_seq = patch_src(batch['src'], opt.src_pad_idx).to(device)
trg_seq, gold = map(lambda x: x.to(device), patch_trg(batch['trg'], opt.trg_pad_idx))
# forward
pred = model(src_seq, trg_seq)
loss, n_correct, n_word = cal_performance(
pred, gold, opt.trg_pad_idx, smoothing=False)
# note keeping
n_word_total += n_word
n_word_correct += n_correct
total_loss += loss.item()
loss_per_word = total_loss/n_word_total
accuracy = n_word_correct/n_word_total
return loss_per_word, accuracy
def train(model, training_data, validation_data, optimizer, device, opt):
''' Start training '''
# Use tensorboard to plot curves, e.g. perplexity, accuracy, learning rate
if opt.use_tb:
print("[Info] Use Tensorboard")
from torch.utils.tensorboard import SummaryWriter
tb_writer = SummaryWriter(log_dir=os.path.join(opt.output_dir, 'tensorboard'))
log_train_file = os.path.join(opt.output_dir, 'train.log')
log_valid_file = os.path.join(opt.output_dir, 'valid.log')
print('[Info] Training performance will be written to file: {} and {}'.format(log_train_file, log_valid_file))
with open(log_train_file, 'w') as log_tf, open(log_valid_file, 'w') as log_vf:
log_tf.write('epoch,loss,ppl,accuracy\n')
log_vf.write('epoch,loss,ppl,accuracy\n')
def print_performances(header, ppl, accu, start_time, lr):
print(' - {header:12} ppl: {ppl: 8.5f}, accuracy: {accu:3.3f} %, lr: {lr:8.5f}, '\
'elapse: {elapse:3.3f} min'.format(
header=f"({header})", ppl=ppl,
accu=100*accu, elapse=(time.time()-start_time)/60, lr=lr))
#valid_accus = []
valid_losses = []
for epoch_i in range(opt.epoch):
print('[ Epoch', epoch_i, ']')
start = time.time()
train_loss, train_accu = train_epoch(model, training_data, optimizer, opt, device, smoothing=opt.label_smoothing)
train_ppl = math.exp(min(train_loss, 100))
# Current learning rate
lr = optimizer._optimizer.param_groups[0]['lr']
print_performances('Training', train_ppl, train_accu, start, lr)
start = time.time()
valid_loss, valid_accu = eval_epoch(model, validation_data, device, opt)
valid_ppl = math.exp(min(valid_loss, 100))
print_performances('Validation', valid_ppl, valid_accu, start, lr)
valid_losses += [valid_loss]
checkpoint = {'epoch': epoch_i, 'model': model.state_dict()}
if opt.save_mode == 'all':
model_name = 'model_accu_{accu:3.3f}.chkpt'.format(accu=100*valid_accu)
torch.save(checkpoint, model_name)
elif opt.save_mode == 'best':
model_name = 'model.chkpt'
if valid_loss <= min(valid_losses):
torch.save(checkpoint, os.path.join(opt.output_dir, model_name))
print(' - [Info] The checkpoint file has been updated.')
with open(log_train_file, 'a') as log_tf, open(log_valid_file, 'a') as log_vf:
log_tf.write('{epoch},{loss: 8.5f},{ppl: 8.5f},{accu:3.3f}\n'.format(
epoch=epoch_i, loss=train_loss,
ppl=train_ppl, accu=100*train_accu))
log_vf.write('{epoch},{loss: 8.5f},{ppl: 8.5f},{accu:3.3f}\n'.format(
epoch=epoch_i, loss=valid_loss,
ppl=valid_ppl, accu=100*valid_accu))
if opt.use_tb:
tb_writer.add_scalars('ppl', {'train': train_ppl, 'val': valid_ppl}, epoch_i)
tb_writer.add_scalars('accuracy', {'train': train_accu*100, 'val': valid_accu*100}, epoch_i)
tb_writer.add_scalar('learning_rate', lr, epoch_i)
def main():
'''
Usage:
python train.py -data_pkl m30k_deen_shr.pkl -log m30k_deen_shr -embs_share_weight -proj_share_weight -label_smoothing -output_dir output -b 256 -warmup 128000
'''
# parser = argparse.ArgumentParser()
# parser.add_argument('-data_pkl', default=None) # all-in-1 data pickle or bpe field
# parser.add_argument('-train_path', default=None) # bpe encoded data
# parser.add_argument('-val_path', default=None) # bpe encoded data
# parser.add_argument('-epoch', type=int, default=10)
# parser.add_argument('-b', '--batch_size', type=int, default=2048)
# parser.add_argument('-d_model', type=int, default=512)
# parser.add_argument('-d_inner_hid', type=int, default=2048)
# parser.add_argument('-d_k', type=int, default=64)
# parser.add_argument('-d_v', type=int, default=64)
# parser.add_argument('-n_head', type=int, default=8)
# parser.add_argument('-n_layers', type=int, default=6)
# parser.add_argument('-warmup','--n_warmup_steps', type=int, default=4000)
# parser.add_argument('-lr_mul', type=float, default=2.0)
# parser.add_argument('-seed', type=int, default=None)
# parser.add_argument('-dropout', type=float, default=0.1)
# parser.add_argument('-embs_share_weight', action='store_true')
# parser.add_argument('-proj_share_weight', action='store_true')
# parser.add_argument('-scale_emb_or_prj', type=str, default='prj')
# parser.add_argument('-output_dir', type=str, default=None)
# parser.add_argument('-use_tb', action='store_true')
# parser.add_argument('-save_mode', type=str, choices=['all', 'best'], default='best')
# parser.add_argument('-no_cuda', action='store_true')
# parser.add_argument('-label_smoothing', action='store_true')
# opt = parser.parse_args()
class Args:
train_src_txt_filepath = './data/dev_zh/newsdev2017-enzh-ref.zh.sgm.txt'
train_tgt_txt_filepaht = './data/dev_en/newsdev2017-zhen-ref.en.sgm.txt'
val_src_txt_filepath = './data/dev_zh/newsdev2017-enzh-ref.zh.sgm.txt'
val_tgt_txt_filepaht = './data/dev_en/newsdev2017-zhen-ref.en.sgm.txt'
seed = 1234
dropout = 0.1
epoch = 100
batch_size = 128
d_model = 512
d_inner_hid = 2048
d_k = 64
d_v = 64
n_head = 8
n_layers = 6
n_warmup_steps = 4000
lr_mul = 2.0
embs_share_weight = False
proj_share_weight = False
scale_emb_or_prj = 'prj'
output_dir = "./output"
save_mode = 'best'
no_cuda = False
label_smoothing = False
use_tb = True
opt = Args()
opt.cuda = not opt.no_cuda
opt.d_word_vec = opt.d_model
# https://pytorch.org/docs/stable/notes/randomness.html
# For reproducibility
if opt.seed is not None:
torch.manual_seed(opt.seed)
torch.backends.cudnn.benchmark = False
# torch.set_deterministic(True)
np.random.seed(opt.seed)
random.seed(opt.seed)
if not opt.output_dir:
print('No experiment result will be saved.')
raise
if not os.path.exists(opt.output_dir):
os.makedirs(opt.output_dir)
if opt.batch_size < 2048 and opt.n_warmup_steps <= 4000:
print('[Warning] The warmup steps may be not enough.\n'\
'(sz_b, warmup) = (2048, 4000) is the official setting.\n'\
'Using smaller batch w/o longer warmup may cause '\
'the warmup stage ends with only little data trained.')
device = torch.device('cuda' if opt.cuda else 'cpu')
#========= Loading Dataset =========#
_, training_data = build_wmt_dataloader(src_txt_filepath=opt.train_src_txt_filepath, tgt_txt_filepath=opt.train_tgt_txt_filepaht, batch_size=opt.batch_size)
_, validation_data = build_wmt_dataloader(src_txt_filepath=opt.train_src_txt_filepath, tgt_txt_filepath=opt.train_tgt_txt_filepaht, batch_size=opt.batch_size)
opt.src_vocab_size = training_data.vocab_size['src']
opt.trg_vocab_size = training_data.vocab_size['tgt']
opt.src_pad_idx = training_data.special_token_index(Constants.PAD_WORD)
opt.trg_pad_idx = training_data.special_token_index(Constants.PAD_WORD)
print(opt)
transformer = Transformer(
opt.src_vocab_size,
opt.trg_vocab_size,
src_pad_idx=opt.src_pad_idx,
trg_pad_idx=opt.trg_pad_idx,
trg_emb_prj_weight_sharing=opt.proj_share_weight,
emb_src_trg_weight_sharing=opt.embs_share_weight,
d_k=opt.d_k,
d_v=opt.d_v,
d_model=opt.d_model,
d_word_vec=opt.d_word_vec,
d_inner=opt.d_inner_hid,
n_layers=opt.n_layers,
n_head=opt.n_head,
dropout=opt.dropout,
scale_emb_or_prj=opt.scale_emb_or_prj).to(device)
optimizer = ScheduledOptim(
optim.Adam(transformer.parameters(), betas=(0.9, 0.98), eps=1e-09),
opt.lr_mul, opt.d_model, opt.n_warmup_steps)
train(transformer, training_data, validation_data, optimizer, device, opt)
if __name__ == '__main__':
main()