def forward(self, x, x_p, mask, mask_p, train=True):
prior, _, attn_dist, _ = self.dec((self.query, x, [], (mask, None)))
prior = self.layer_norm1(prior)
mean = self.mean(prior)
log_var = self.var(prior)
eps = torch.randn(prior.size())
std = torch.exp(0.5 * log_var)
if config.USE_CUDA: eps = eps.cuda()
z = eps * std + mean
kld_loss = 0
if x_p is not None:
posterior, _, attn_dist_p, _ = self.var_dec(
(self.query, x_p, [], (mask_p, None)))
posterior = self.layer_norm2(posterior)
mean_p = self.mean_p(posterior)
log_var_p = self.var_p(posterior)
kld_loss = gaussian_kld(mean_p, log_var_p, mean, log_var)
kld_loss = torch.mean(kld_loss)
if train:
std = torch.exp(0.5 * log_var_p)
if config.USE_CUDA: eps = eps.cuda()
z = eps * std + mean_p
return kld_loss, z
def train_one_batch(self, batch, iter, train=True):
enc_batch, _, _, enc_batch_extend_vocab, extra_zeros, _, _ = get_input_from_batch(
batch)
dec_batch, _, _, _, _ = get_output_from_batch(batch)
if (config.noam):
self.optimizer.optimizer.zero_grad()
else:
self.optimizer.zero_grad()
## Response encode
mask_res = batch["posterior_batch"].data.eq(
config.PAD_idx).unsqueeze(1)
posterior_mask = self.embedding(batch["posterior_mask"])
r_encoder_outputs = self.r_encoder(
self.embedding(batch["posterior_batch"]), mask_res)
## Encode
mask_src = enc_batch.data.eq(config.PAD_idx).unsqueeze(1)
emb_mask = self.embedding(batch["input_mask"])
encoder_outputs = self.encoder(self.embedding(enc_batch), mask_src)
meta = self.embedding(batch["program_label"])
# Decode
mask_trg = dec_batch.data.eq(config.PAD_idx).unsqueeze(1)
latent_dim = meta.size()[-1]
meta = meta.repeat(1, dec_batch.size(1)).view(dec_batch.size(0),
dec_batch.size(1),
latent_dim)
pre_logit, attn_dist, mean, log_var = self.decoder(
meta, encoder_outputs, r_encoder_outputs,
(mask_src, mask_res, mask_trg))
if not train:
pre_logit, attn_dist, _, _ = self.decoder(
meta, encoder_outputs, None, (mask_src, None, mask_trg))
## compute output dist
logit = self.generator(
pre_logit,
attn_dist,
enc_batch_extend_vocab if config.pointer_gen else None,
extra_zeros,
attn_dist_db=None)
## loss: NNL if ptr else Cross entropy
loss_rec = self.criterion(logit.contiguous().view(-1, logit.size(-1)),
dec_batch.contiguous().view(-1))
kld_loss = gaussian_kld(mean["posterior"], log_var["posterior"],
mean["prior"], log_var["prior"])
kld_loss = torch.mean(kld_loss)
kl_weight = min(iter / config.full_kl_step,
1) if config.full_kl_step > 0 else 1.0
loss = loss_rec + config.kl_ceiling * kl_weight * kld_loss
if (train):
loss.backward()
# clip gradient
nn.utils.clip_grad_norm_(self.parameters(), config.max_grad_norm)
self.optimizer.step()
return loss_rec.item(), math.exp(min(loss_rec.item(),
100)), kld_loss.item()
def train_n_batch(self, batchs, iter, train=True):
if(config.noam):
self.optimizer.optimizer.zero_grad()
else:
self.optimizer.zero_grad()
for batch in batchs:
enc_batch, _, _, enc_batch_extend_vocab, extra_zeros, _, _ = get_input_from_batch(batch)
dec_batch, _, _, _, _ = get_output_from_batch(batch)
## Encode
mask_src = enc_batch.data.eq(config.PAD_idx).unsqueeze(1)
encoder_outputs = self.encoder(self.embedding(enc_batch), mask_src)
meta = self.embedding(batch["program_label"])
if config.dataset=="empathetic":
meta = meta-meta
# Decode
sos_token = torch.LongTensor([config.SOS_idx] * enc_batch.size(0)).unsqueeze(1)
if config.USE_CUDA: sos_token = sos_token.cuda()
dec_batch_shift = torch.cat((sos_token,dec_batch[:, :-1]),1)
mask_trg = dec_batch_shift.data.eq(config.PAD_idx).unsqueeze(1)
pre_logit, attn_dist, mean, log_var, probs= self.decoder(self.embedding(dec_batch_shift)+meta.unsqueeze(1),encoder_outputs, True, (mask_src,mask_trg))
## compute output dist
logit = self.generator(pre_logit,attn_dist,enc_batch_extend_vocab if config.pointer_gen else None, extra_zeros, attn_dist_db=None)
## loss: NNL if ptr else Cross entropy
sbow = dec_batch #[batch, seq_len]
seq_len = sbow.size(1)
loss_rec = self.criterion(logit.contiguous().view(-1, logit.size(-1)), dec_batch.contiguous().view(-1))
if config.model=="cvaetrs":
loss_aux = 0
for prob in probs:
sbow_mask = _get_attn_subsequent_mask(seq_len).transpose(1,2)
sbow.unsqueeze(2).repeat(1,1,seq_len).masked_fill_(sbow_mask,config.PAD_idx)#[batch, seq_len, seq_len]
loss_aux+= self.criterion(prob.contiguous().view(-1, prob.size(-1)), sbow.contiguous().view(-1))
kld_loss = gaussian_kld(mean["posterior"], log_var["posterior"],mean["prior"], log_var["prior"])
kld_loss = torch.mean(kld_loss)
kl_weight = min(math.tanh(6 * iter/config.full_kl_step - 3) + 1, 1)
#kl_weight = min(iter/config.full_kl_step, 1) if config.full_kl_step >0 else 1.0
loss = loss_rec + config.kl_ceiling * kl_weight*kld_loss + config.aux_ceiling*loss_aux
elbo = loss_rec+kld_loss
else:
loss = loss_rec
elbo = loss_rec
kld_loss = torch.Tensor([0])
loss_aux = torch.Tensor([0])
loss.backward()
# clip gradient
nn.utils.clip_grad_norm_(self.parameters(), config.max_grad_norm)
self.optimizer.step()
return loss_rec.item(), math.exp(min(loss_rec.item(), 100)), kld_loss.item(), loss_aux.item(), elbo.item()
def forward(self, x, x_p, train=True):
mean = self.mean(x)
log_var = self.var(x)
eps = torch.randn(mean.size())
std = torch.exp(0.5 * log_var)
if config.USE_CUDA: eps = eps.cuda()
z = eps * std + mean
kld_loss = 0
if x_p is not None:
mean_p = self.mean_p(torch.cat((x_p, x), dim=-1))
log_var_p = self.var_p(torch.cat((x_p, x), dim=-1))
kld_loss = gaussian_kld(mean_p, log_var_p, mean, log_var)
kld_loss = torch.mean(kld_loss)
if train:
std = torch.exp(0.5 * log_var_p)
if config.USE_CUDA: eps = eps.cuda()
z = eps * std + mean_p
return kld_loss, z
def train_one_batch(self, batch, iter, train=True):
enc_batch, enc_padding_mask, enc_lens, enc_batch_extend_vocab, extra_zeros, _, _ = get_input_from_batch(batch)
dec_batch, _, _, _, _ = get_output_from_batch(batch)
if(config.noam):
self.optimizer.optimizer.zero_grad()
else:
self.optimizer.zero_grad()
## Response encode
mask_res = batch["posterior_batch"].data.eq(config.PAD_idx).unsqueeze(1)
post_emb = self.embedding(batch["posterior_batch"])
r_encoder_outputs = self.r_encoder(post_emb, mask_res)
## Encode
num_sentences, enc_seq_len = enc_batch.size()
batch_size = enc_lens.size(0)
max_len = enc_lens.data.max().item()
input_lengths = torch.sum(~enc_batch.data.eq(config.PAD_idx), dim=1)
# word level encoder
enc_emb = self.embedding(enc_batch)
word_encoder_outpus, word_encoder_hidden = self.word_encoder(enc_emb, input_lengths)
word_encoder_hidden = word_encoder_hidden.transpose(1, 0).reshape(num_sentences, -1)
# pad and pack word_encoder_hidden
start = torch.cumsum(torch.cat((enc_lens.data.new(1).zero_(), enc_lens[:-1])), 0)
word_encoder_hidden = torch.stack([pad(word_encoder_hidden.narrow(0, s, l), max_len)
for s, l in zip(start.data.tolist(), enc_lens.data.tolist())], 0)
# mask_src = ~(enc_padding_mask.bool()).unsqueeze(1)
mask_src = (1 - enc_padding_mask.byte()).unsqueeze(1)
# context level encoder
if word_encoder_hidden.size(-1) != config.hidden_dim:
word_encoder_hidden = self.linear(word_encoder_hidden)
encoder_outputs = self.encoder(word_encoder_hidden, mask_src)
# Decode
sos_token = torch.LongTensor([config.SOS_idx] * batch_size).unsqueeze(1)
if config.USE_CUDA: sos_token = sos_token.cuda()
dec_batch_shift = torch.cat((sos_token, dec_batch[:, :-1]), 1) #(batch, len, embedding)
mask_trg = dec_batch_shift.data.eq(config.PAD_idx).unsqueeze(1)
dec_emb = self.embedding(dec_batch_shift)
pre_logit, attn_dist, mean, log_var, probs = self.decoder(dec_emb, encoder_outputs, r_encoder_outputs,
(mask_src, mask_res, mask_trg))
## compute output dist
logit = self.generator(pre_logit, attn_dist, enc_batch_extend_vocab if config.pointer_gen else None, extra_zeros, attn_dist_db=None)
## loss: NNL if ptr else Cross entropy
sbow = dec_batch #[batch, seq_len]
seq_len = sbow.size(1)
loss_rec = self.criterion(logit.contiguous().view(-1, logit.size(-1)), dec_batch.contiguous().view(-1))
if config.model=="cvaetrs":
loss_aux = 0
for prob in probs:
sbow_mask = _get_attn_subsequent_mask(seq_len).transpose(1,2)
sbow.unsqueeze(2).repeat(1,1,seq_len).masked_fill_(sbow_mask,config.PAD_idx)#[batch, seq_len, seq_len]
loss_aux+= self.criterion(prob.contiguous().view(-1, prob.size(-1)), sbow.contiguous().view(-1))
kld_loss = gaussian_kld(mean["posterior"], log_var["posterior"],mean["prior"], log_var["prior"])
kld_loss = torch.mean(kld_loss)
kl_weight = min(math.tanh(6 * iter/config.full_kl_step - 3) + 1, 1)
#kl_weight = min(iter/config.full_kl_step, 1) if config.full_kl_step >0 else 1.0
loss = loss_rec + config.kl_ceiling * kl_weight*kld_loss + config.aux_ceiling*loss_aux
elbo = loss_rec + kld_loss
else:
loss = loss_rec
elbo = loss_rec
kld_loss = torch.Tensor([0])
loss_aux = torch.Tensor([0])
if(train):
loss.backward()
# clip gradient
nn.utils.clip_grad_norm_(self.parameters(), config.max_grad_norm)
self.optimizer.step()
return loss_rec.item(), math.exp(min(loss_rec.item(), 100)), kld_loss.item(), loss_aux.item(), elbo.item()