-
Notifications
You must be signed in to change notification settings - Fork 0
/
main_flow.py
270 lines (203 loc) · 10.3 KB
/
main_flow.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
import os
import sys
import shutil
from tqdm import tqdm, tqdm_notebook
from math import log, sqrt, pi
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
import numpy as np
from config import *
sys.path.append('../')
# from cls import *
from loss import *
from utils import *
from dataset import TVQADataset, pad_collate, preprocess_inputs
def calc_loss(log_p, logdet, image_size, n_bins, input_hidden=64):
n_pixel = input_hidden
# image_size *
loss = -log(n_bins) * n_pixel
loss = loss + logdet + log_p
return (
(-loss / (log(2) * n_pixel)).mean(),
(log_p / (log(2) * n_pixel)).mean(),
(logdet / (log(2) * n_pixel)).mean(),
)
def line_to_words(line):
words = line.lower().split()
words = [w for w in words if w != ","]
words = [w if w != '<eos>' else '[SEP]' for w in words ]
return words
def line_to_words_vcpt(vcpt_sentence):
attr_obj_pairs = vcpt_sentence.lower().split(",") # comma is also removed
unique_pairs = []
for pair in attr_obj_pairs:
if pair not in unique_pairs:
unique_pairs.append(pair)
words = []
for pair in unique_pairs:
words.extend(pair.split())
return words
def train(opt, dset, model, flow, criterions, optimizer, epoch, previous_best_loss, schedular, tokenizer=None, bert=None, container1=None, container2=None, z_sample=None, cos=None):
dset.set_mode('train')
model.train()
flow.train()
train_loader = DataLoader(dset, batch_size=opt.bsz, shuffle=False, collate_fn=pad_collate)
train_loss = []
train_nll_loss = []
train_recon_loss = []
valid_loss_log = ["batch_idx\tloss\trecon-loss\tnll-loss\tvalid-loss\tvalid-recon-loss\tvalid-nll-loss"]
torch.set_grad_enabled(True)
beam = BeamSearcher()
for batch_idx, batch in tqdm(enumerate(train_loader)):
model_inputs, labels, qids = preprocess_inputs(batch, opt.max_sub_l, opt.max_vcpt_l, opt.max_vid_l, device=opt.device)
model_inputs = [flow] + model_inputs
flow_inputs, inputs_length = model.get_conditional(*model_inputs)
log_p, logdet, flow_z, inputs = model.encode(flow, flow_inputs, inputs_length)
recon_loss = 0
nll_loss, log_p, log_det = calc_loss(log_p, logdet.mean(), inputs_length.float().mean(), 64, 300)
loss = nll_loss
optimizer[0].zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(flow.parameters(), 1.0)
optimizer[0].step()
train_loss+=[loss.item()]
train_nll_loss+=[nll_loss.item()]
train_recon_loss+=[recon_loss]
if batch_idx % opt.log_freq == 0:
niter = epoch * len(train_loader) + batch_idx
opt.writer.add_scalar("Train/Loss", np.mean(train_loss[-opt.log_freq:]), niter)
opt.writer.add_scalar("Train/NLL-Loss", np.mean(train_nll_loss[-opt.log_freq:]), niter)
opt.writer.add_scalar("Train/Recon-Loss", np.mean(train_recon_loss[-opt.log_freq:]), niter)
valid_loss, valid_nll_loss, valid_recon_loss, val_logp, sentenses, recon_error = validate(opt, dset, model, flow, beam, mode="valid", cos=cos)
opt.writer.add_scalar("Valid/Loss", valid_loss, niter)
opt.writer.add_scalar("Valid/NLL-Loss", valid_nll_loss, niter)
opt.writer.add_scalar("Valid/Recon-Loss", valid_recon_loss, niter)
valid_log_str ="%02d\t%.4f\t%.4f\t%.4f\t%.4f\t%.4f\t%.4f" % (batch_idx, loss, recon_loss, nll_loss, valid_loss, valid_recon_loss, valid_nll_loss)
valid_loss_log.append(valid_log_str)
if valid_loss < previous_best_loss:
previous_best_loss = valid_loss
torch.save(flow.state_dict(), os.path.join(opt.results_dir, "best_valid_flow.pth"))
torch.save(optimizer[0].state_dict(), os.path.join(opt.results_dir, "optimizer1.pth"))
print(" Train Epoch %d loss %.4f nll-loss %.4f recon-loss %.4f Val loss %.4f nll-loss %.4f recon-loss %.4f logp %.4f"
% (epoch, np.mean(train_loss[-opt.log_freq:]), np.mean(train_nll_loss[-opt.log_freq:]), np.mean(train_recon_loss[-opt.log_freq:]), valid_loss, valid_nll_loss, valid_recon_loss, val_logp))
for i in range(len(sentenses)):
print(sentenses[i])
torch.set_grad_enabled(True)
model.train()
flow.train()
dset.set_mode("train")
if opt.debug:
break
with open(os.path.join(opt.results_dir, "valid_loss.log"), "a") as f:
f.write("\n".join(valid_loss_log) + "\n")
return previous_best_loss
def validate(opt, dset, model, flow, beam, mode="valid", cos=None):
dset.set_mode(mode)
torch.set_grad_enabled(False)
model.eval()
flow.eval()
valid_loader = DataLoader(dset, batch_size=opt.test_bsz, shuffle=False, collate_fn=pad_collate)
valid_loss = []
valid_nll_loss = []
valid_recon_loss = []
val_logp = []
for k, batch in enumerate(valid_loader):
model_inputs, labels, qids = preprocess_inputs(batch, opt.max_sub_l, opt.max_vcpt_l, opt.max_vid_l, device=opt.device)
model_inputs = [flow] + model_inputs
flow_inputs, inputs_length = model.get_conditional(*model_inputs)
log_p, logdet, flow_z, inputs = model.encode(flow, flow_inputs, inputs_length)
if k == 0:
reconstructed, logits = model.decode(flow, inputs, flow_z, True, 1, cos)
recon_error = torch.abs(reconstructed[0][:model_inputs[2][0]] - inputs[0][:model_inputs[2][0]]).mean()
print('reconstruction error:', recon_error)
if recon_error > 2.0:
print('----warning, invertibility_test failed----')
print(logits[0].argmax(-1), model_inputs[1][0][1:])
recon_loss = 0
nll_loss, log_p, log_det = calc_loss(log_p, logdet.mean(), model_inputs[12].float().mean(), 1, 300)
loss = nll_loss
val_logp+=[log_p.item()]
valid_loss+=[loss.item()]
valid_nll_loss+=[nll_loss.item()]
valid_recon_loss+=[recon_loss]
if k == 0:
show_size = 5
if opt.use_ar:
z_sample = model.get_sample(flow_z, show_size=show_size)
sentenses = []
for i in range(show_size):
sentenses.append(beam.beam_search(flow, flow_inputs[:1, :2], inputs_length[:1], [z_sample[0][i:i+1]], dset, i, beam_size=opt.beam_size, return_tokens=False, max_decode_steps=model.max_len-1, early_stop=False, cos=cos))
sentenses_greedy = []
for i in range(show_size):
sentenses_greedy.append(beam.beam_search(flow, flow_inputs[:1, :2], inputs_length[:1], [z_sample[0][i:i+1]], dset, i, beam_size=1, return_tokens=False, max_decode_steps=model.max_len-1, early_stop=False, cos=cos))
else:
reconstructed, logits = model.decode(flow, inputs, flow_z, False, show_size, cos)
sentenses = get_sentences(model, logits.view(-1, logits.size(-1)), dset, model.max_len, show_size, early_stop=False, show_all=True)
if opt.debug:
break
if opt.val_steps:
if opt.val_steps <= k+1:
break
valid_loss = np.mean(valid_loss)
valid_nll_loss = np.mean(valid_nll_loss)
valid_recon_loss = np.mean(valid_recon_loss)
val_logp = np.mean(val_logp)
return valid_loss, valid_nll_loss, valid_recon_loss, val_logp, sentenses, recon_error
if __name__ == "__main__":
save_numpy = []
save_numpy_index = []
torch.manual_seed(2020)
opt = BaseOptions().parse()
cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"]=str(opt.gpus[0])
from model.flow_entry import LangFlow
if opt.use_ar:
from get_tokens_ar import *
from model.flow.flowauto_eb4 import Glow
else:
from get_tokens import *
from model.flow.flownonauto import Glow
writer = SummaryWriter(opt.results_dir)
opt.writer = writer
z_sample = None
dset = TVQADataset(opt)
opt.vocab_size = len(dset.word2idx)
model = LangFlow(opt)
model.load_embedding(dset.vocab_embedding)
model.cuda()
flow = Glow(opt, embedding=model.embedding).cuda()
cos = nn.CosineSimilarity(dim=2, eps=1e-12)
flow.embedding = model.embedding
flow.embedding.weight.requires_grad = False
if opt.restore_name:
flow.load_state_dict(torch.load(opt.results_dir_base+opt.restore_name+'/best_valid_flow.pth', map_location='cuda:0'), strict=False)
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)/1000000.0
print('The number of parameters of model is', num_params, "M")
num_params = sum(p.numel() for p in flow.parameters() if p.requires_grad)/1000000.0
print('The number of parameters of flow is', num_params, "M")
criterions = [EntropyLoss().cuda(), NLL().cuda()]
optimizer = [torch.optim.Adam(filter(lambda p: p.requires_grad, flow.parameters()), lr=opt.lr)]
if opt.restore_name:
di=opt.results_dir_base+opt.restore_name+'/optimizer1.pth'
optimizer[0].load_state_dict(torch.load(di, map_location='cuda:0'))
schedular = None
best_loss = 1e10
early_stopping_cnt = 0
early_stopping_flag = False
for epoch in range(opt.n_epoch):
if not early_stopping_flag:
# train for one epoch, valid per n batches, save the log and the best model
cur_loss = train(opt, dset, model, flow, criterions, optimizer, epoch, best_loss, schedular, container1=save_numpy, container2=save_numpy_index, z_sample=z_sample, cos=cos)
# remember best acc
best_loss = min(cur_loss, best_loss)
else:
print("early stop with valid loss %.4f" % best_loss)
opt.writer.export_scalars_to_json(os.path.join(opt.results_dir, "all_scalars.json"))
opt.writer.close()
break # early stop break
if opt.debug:
break