class PoemWAE(nn.Module): def __init__(self, config, api, PAD_token=0, pretrain_weight=None): super(PoemWAE, self).__init__() self.vocab = api.vocab self.vocab_size = len(self.vocab) self.rev_vocab = api.rev_vocab self.go_id = self.rev_vocab["<s>"] self.eos_id = self.rev_vocab["</s>"] self.maxlen = config.maxlen self.clip = config.clip self.lambda_gp = config.lambda_gp self.lr_gan_g = config.lr_gan_g self.lr_gan_d = config.lr_gan_d self.n_d_loss = config.n_d_loss self.temp = config.temp self.init_w = config.init_weight self.embedder = nn.Embedding(self.vocab_size, config.emb_size, padding_idx=PAD_token) if pretrain_weight is not None: self.embedder.weight.data.copy_(torch.from_numpy(pretrain_weight)) # 用同一个seq_encoder来编码标题和前后两句话 self.seq_encoder = Encoder(self.embedder, config.emb_size, config.n_hidden, True, config.n_layers, config.noise_radius) # 由于Poem这里context是title和last sentence双向GRU编码后的直接cat,4*hidden # 注意如果使用Poemwar_gmp则使用子类中的prior_net,即混合高斯分布的一个先验分布 self.prior_net = Variation(config.n_hidden * 4, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # p(e|c) # 注意这儿原来是给Dialog那个任务用的,3*hidden # Poem数据集上,将title和上一句,另外加上x都分别用双向GRU编码并cat,因此是6*hidden self.post_net = Variation(config.n_hidden * 6, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) self.post_generator = nn.Sequential( nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size)) self.post_generator.apply(self.init_weights) self.prior_generator = nn.Sequential( nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size)) self.prior_generator.apply(self.init_weights) self.init_decoder_hidden = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden * 4), nn.BatchNorm1d(config.n_hidden * 4, eps=1e-05, momentum=0.1), nn.ReLU()) # 由于Poem这里context是title和last sentence双向GRU编码后的直接cat,因此hidden_size变为z_size + 4*hidden # 修改:decoder的hidden_size还设为n_hidden, init_hidden使用一个MLP将cat变换为n_hidden self.decoder = Decoder(self.embedder, config.emb_size, config.n_hidden * 4, self.vocab_size, n_layers=1) self.discriminator = nn.Sequential( # 因为Poem的cat两个双向编码,这里改为4*n_hidden + z_size nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden * 2), nn.BatchNorm1d(config.n_hidden * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config.n_hidden * 2, config.n_hidden * 2), nn.BatchNorm1d(config.n_hidden * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config.n_hidden * 2, 1), ) self.discriminator.apply(self.init_weights) # optimizer 定义,分别对应三个模块的训练,注意!三个模块的optimizer不相同 # self.optimizer_AE = optim.SGD(list(self.seq_encoder.parameters()) self.optimizer_AE = optim.SGD( list(self.seq_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.init_decoder_hidden.parameters()) + list(self.decoder.parameters()), lr=config.lr_ae) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()), lr=self.lr_gan_g) self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=self.lr_gan_d) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.8) self.criterion_ce = nn.CrossEntropyLoss() def init_weights(self, m): if isinstance(m, nn.Linear): m.weight.data.uniform_(-self.init_w, self.init_w) # nn.init.kaiming_normal_(m.weight.data) # nn.init.kaiming_uniform_(m.weight.data) m.bias.data.fill_(0) # x: (batch, 2*n_hidden) # c: (batch, 2*2*n_hidden) def sample_code_post(self, x, c): z, _, _ = self.post_net(torch.cat((x, c), 1)) # 输入:(batch, 3*2*n_hidden) z = self.post_generator(z) return z def sample_code_prior_sentiment(self, c, align): choice_statistic = self.prior_net(c, align) # e: (batch, z_size) return choice_statistic def sample_code_prior(self, c): z, _, _ = self.prior_net(c) # e: (batch, z_size) z = self.prior_generator(z) # z: (batch, z_size) return z # 输入 title, context, target, target_lens. # c由title和context encode之后的hidden相concat而成 def train_AE(self, title, context, target, target_lens): self.seq_encoder.train() self.decoder.train() # import pdb # pdb.set_trace() # (batch, 2 * hidden_size) title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) # (batch, 2 * hidden_size) x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) # context_embedding c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 2 * hidden_size * 2) z = self.sample_code_post(x, c) # (batch, z_size) # 标准的autoencoder的decode,decoder初态为x, c的cat,将target错位输入 # output: (batch, len, vocab_size) len是9,即7+标点+</s> output = self.decoder(self.init_decoder_hidden(torch.cat((z, c), 1)), None, target[:, :-1], target_lens - 1) flattened_output = output.view(-1, self.vocab_size) dec_target = target[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # 即判断target的token中是否有0(pad项) masked_target = dec_target.masked_select(mask) # 选出非pad项 output_mask = mask.unsqueeze(1).expand( mask.size(0), self.vocab_size) # [(batch_sz * seq_len) x n_tokens] masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) self.optimizer_AE.zero_grad() loss = self.criterion_ce(masked_output / self.temp, masked_target) loss.backward() torch.nn.utils.clip_grad_norm_( list(self.seq_encoder.parameters()) + list(self.decoder.parameters()), self.clip) self.optimizer_AE.step() return [('train_loss_AE', loss.item())] # G是来缩短W距离的,可以类比VAE里面的缩小KL散度项 def train_G(self, title, context, target, target_lens, sentiment_mask=None, mask_type=None): self.seq_encoder.eval() self.optimizer_G.zero_grad() for p in self.discriminator.parameters(): p.requires_grad = False title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 2 * hidden_size * 2) # -----------------posterior samples --------------------------- x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) z_post = self.sample_code_post( x.detach(), c.detach()) # 去掉梯度,防止梯度向encoder的传播 (batch, z_size) errG_post = torch.mean( self.discriminator(torch.cat( (z_post, c.detach()), 1))) * self.n_d_loss # (batch, z_size + 4 * hidden) errG_post.backward(minus_one) # ----------------- prior samples --------------------------- prior_z = self.sample_code_prior(c.detach()) errG_prior = torch.mean( self.discriminator(torch.cat( (prior_z, c.detach()), 1))) * self.n_d_loss # import pdb # pdb.set_trace() errG_prior.backward(one) self.optimizer_G.step() for p in self.discriminator.parameters(): p.requires_grad = True costG = errG_prior - errG_post return [('train_loss_G', costG.item())] # D是用来拟合W距离,loss下降说明拟合度变好,增大gradient_penalty一定程度上可以提高拟合度 # n_iters_n越大,D训练的次数越多,对应的拟合度也越好 def train_D(self, title, context, target, target_lens): self.seq_encoder.eval() self.discriminator.train() self.optimizer_D.zero_grad() batch_size = context.size(0) title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 2, hidden_size * 2) x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) post_z = self.sample_code_post(x, c) errD_post = torch.mean( self.discriminator(torch.cat( (post_z.detach(), c.detach()), 1))) * self.n_d_loss errD_post.backward(one) prior_z = self.sample_code_prior(c) errD_prior = torch.mean( self.discriminator(torch.cat( (prior_z.detach(), c.detach()), 1))) * self.n_d_loss errD_prior.backward(minus_one) # import pdb # pdb.set_trace() alpha = to_tensor(torch.rand(batch_size, 1)) alpha = alpha.expand(prior_z.size()) interpolates = alpha * prior_z.data + ((1 - alpha) * post_z.data) interpolates = Variable(interpolates, requires_grad=True) d_input = torch.cat((interpolates, c.detach()), 1) disc_interpolates = torch.mean(self.discriminator(d_input)) gradients = torch.autograd.grad( outputs=disc_interpolates, inputs=interpolates, grad_outputs=to_tensor(torch.ones(disc_interpolates.size())), create_graph=True, retain_graph=True, only_inputs=True)[0] gradient_penalty = ( (gradients.contiguous().view(gradients.size(0), -1).norm(2, dim=1) - 1)**2).mean() * self.lambda_gp gradient_penalty.backward() self.optimizer_D.step() costD = -(errD_prior - errD_post) + gradient_penalty return [('train_loss_D', costD.item())] def valid(self, title, context, target, target_lens, sentiment_mask=None): self.seq_encoder.eval() self.discriminator.eval() self.decoder.eval() title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 2 * hidden_size * 2) x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) post_z = self.sample_code_post(x, c) prior_z = self.sample_code_prior(c) errD_post = torch.mean(self.discriminator(torch.cat((post_z, c), 1))) errD_prior = torch.mean(self.discriminator(torch.cat((prior_z, c), 1))) costD = -(errD_prior - errD_post) costG = -costD dec_target = target[:, 1:].contiguous().view(-1) # (batch_size * len) mask = dec_target.gt(0) # 即判断target的token中是否有0(pad项) masked_target = dec_target.masked_select(mask) # 选出非pad项 output_mask = mask.unsqueeze(1).expand(mask.size(0), self.vocab_size) output = self.decoder( self.init_decoder_hidden(torch.cat((post_z, c), 1)), None, target[:, :-1], (target_lens - 1)) flattened_output = output.view(-1, self.vocab_size) masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) lossAE = self.criterion_ce(masked_output / self.temp, masked_target) return [('valid_loss_AE', lossAE.item()), ('valid_loss_G', costG.item()), ('valid_loss_D', costD.item())] # 正如论文中说的,测试生成的时候,从先验网络中拿到噪声,用G生成prior_z(即代码中的sample_code_prior(c)) # 然后decoder将prior_z和c的cat当做输入,decode出这句诗(这和论文里面不太一样,论文里面只把prior_z当做输入) # batch_size是1,一次测一句 # title 即标题 # context 上一句 def test(self, title_tensor, title_words, headers): self.seq_encoder.eval() self.discriminator.eval() self.decoder.eval() # tem初始化为[2,3,0,0,0,0,0,0,0] tem = [[2, 3] + [0] * (self.maxlen - 2)] pred_poems = [] title_tokens = [ self.vocab[e] for e in title_words[0].tolist() if e not in [0, self.eos_id, self.go_id] ] pred_poems.append(title_tokens) for sent_id in range(4): tem = to_tensor(np.array(tem)) context = tem # vec_context = np.zeros((batch_size, self.maxlen), dtype=np.int64) # for b_id in range(batch_size): # vec_context[b_id, :] = np.array(context[b_id]) # context = to_tensor(vec_context) title_last_hidden, _ = self.seq_encoder( title_tensor) # (batch=1, 2*hidden) if sent_id == 0: context_last_hidden, _ = self.seq_encoder( title_tensor) # (batch=1, 2*hidden) else: context_last_hidden, _ = self.seq_encoder( context) # (batch=1, 2*hidden) c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 4*hidden_size) # 由于一次只有一首诗,batch_size = 1,因此不必repeat prior_z = self.sample_code_prior(c) # decode_words 是完整的一句诗 decode_words = self.decoder.testing( init_hidden=self.init_decoder_hidden(torch.cat((prior_z, c), 1)), maxlen=self.maxlen, go_id=self.go_id, mode="greedy", header=headers[sent_id]) decode_words = decode_words[0].tolist() # import pdb # pdb.set_trace() if len(decode_words) > self.maxlen: tem = [decode_words[0:self.maxlen]] else: tem = [[0] * (self.maxlen - len(decode_words)) + decode_words] pred_tokens = [ self.vocab[e] for e in decode_words[:-1] if e != self.eos_id and e != 0 ] pred_poems.append(pred_tokens) gen = '' for line in pred_poems: true_str = " ".join(line) gen = gen + true_str + '\n' return gen def sample(self, title, context, repeat, go_id, end_id): self.seq_encoder.eval() self.decoder.eval() title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) c = torch.cat((title_last_hidden, context_last_hidden), 1) # (batch, 2 * hidden_size * 2) c_repeated = c.expand( repeat, -1) # 注意,我们输入的batch_size是1,这里复制repeat遍,为了后面的BLEU计算 prior_z = self.sample_code_prior( c_repeated) # c_repeated: (batch_size=repeat, 4*hidden_size) # (batch, max_len, 1) (batch_size, 1) sample_words, sample_lens = self.decoder.sampling( self.init_decoder_hidden(torch.cat((prior_z, c_repeated), 1)), self.maxlen, go_id, end_id, "greedy") return sample_words, sample_lens
def __init__(self, config, api, PAD_token=0, pretrain_weight=None): super(PoemWAE, self).__init__() self.vocab = api.vocab self.vocab_size = len(self.vocab) self.rev_vocab = api.rev_vocab self.go_id = self.rev_vocab["<s>"] self.eos_id = self.rev_vocab["</s>"] self.maxlen = config.maxlen self.clip = config.clip self.lambda_gp = config.lambda_gp self.lr_gan_g = config.lr_gan_g self.lr_gan_d = config.lr_gan_d self.n_d_loss = config.n_d_loss self.temp = config.temp self.init_w = config.init_weight self.embedder = nn.Embedding(self.vocab_size, config.emb_size, padding_idx=PAD_token) if pretrain_weight is not None: self.embedder.weight.data.copy_(torch.from_numpy(pretrain_weight)) # 用同一个seq_encoder来编码标题和前后两句话 self.seq_encoder = Encoder(self.embedder, config.emb_size, config.n_hidden, True, config.n_layers, config.noise_radius) # 由于Poem这里context是title和last sentence双向GRU编码后的直接cat,4*hidden # 注意如果使用Poemwar_gmp则使用子类中的prior_net,即混合高斯分布的一个先验分布 self.prior_net = Variation(config.n_hidden * 4, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # p(e|c) # 注意这儿原来是给Dialog那个任务用的,3*hidden # Poem数据集上,将title和上一句,另外加上x都分别用双向GRU编码并cat,因此是6*hidden self.post_net = Variation(config.n_hidden * 6, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) self.post_generator = nn.Sequential( nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size)) self.post_generator.apply(self.init_weights) self.prior_generator = nn.Sequential( nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size), nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config.z_size, config.z_size)) self.prior_generator.apply(self.init_weights) self.init_decoder_hidden = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden * 4), nn.BatchNorm1d(config.n_hidden * 4, eps=1e-05, momentum=0.1), nn.ReLU()) # 由于Poem这里context是title和last sentence双向GRU编码后的直接cat,因此hidden_size变为z_size + 4*hidden # 修改:decoder的hidden_size还设为n_hidden, init_hidden使用一个MLP将cat变换为n_hidden self.decoder = Decoder(self.embedder, config.emb_size, config.n_hidden * 4, self.vocab_size, n_layers=1) self.discriminator = nn.Sequential( # 因为Poem的cat两个双向编码,这里改为4*n_hidden + z_size nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden * 2), nn.BatchNorm1d(config.n_hidden * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config.n_hidden * 2, config.n_hidden * 2), nn.BatchNorm1d(config.n_hidden * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config.n_hidden * 2, 1), ) self.discriminator.apply(self.init_weights) # optimizer 定义,分别对应三个模块的训练,注意!三个模块的optimizer不相同 # self.optimizer_AE = optim.SGD(list(self.seq_encoder.parameters()) self.optimizer_AE = optim.SGD( list(self.seq_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.init_decoder_hidden.parameters()) + list(self.decoder.parameters()), lr=config.lr_ae) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()), lr=self.lr_gan_g) self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=self.lr_gan_d) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.8) self.criterion_ce = nn.CrossEntropyLoss()
class DFVAE(nn.Module): def __init__(self, config, vocab_size, PAD_token=0): super(DFVAE, self).__init__() self.vocab_size = vocab_size self.maxlen = config['maxlen'] self.clip = config['clip'] self.lambda_gp = config['lambda_gp'] self.temp = config['temp'] self.embedder = nn.Embedding(vocab_size, config['emb_size'], padding_idx=PAD_token) self.utt_encoder = Encoder(self.embedder, config['emb_size'], config['n_hidden'], True, config['n_layers'], config['noise_radius']) self.context_encoder = ContextEncoder(self.utt_encoder, config['n_hidden'] * 2 + 2, config['n_hidden'], 1, config['noise_radius']) self.prior_net = Variation(config['n_hidden'], config['z_size']) # p(e|c) self.post_net = Variation(config['n_hidden'] * 3, config['z_size']) # q(e|c,x) #self.prior_highway = nn.Linear(config['n_hidden'], config['n_hidden']) #self.post_highway = nn.Linear(config['n_hidden'] * 3, config['n_hidden']) self.postflow1 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.postflow2 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.postflow3 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow1 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow2 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow3 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.post_generator = nn_.SequentialFlow(self.postflow1, self.postflow2, self.postflow3) self.prior_generator = nn_.SequentialFlow(self.priorflow1, self.priorflow2, self.priorflow3) self.decoder = Decoder(self.embedder, config['emb_size'], config['n_hidden'] + config['z_size'], vocab_size, n_layers=1) self.optimizer_AE = optim.SGD( list(self.context_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.decoder.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()) #+list(self.prior_highway.parameters()) #+list(self.post_highway.parameters()) , lr=config['lr_ae']) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()) #+list(self.prior_highway.parameters()) #+list(self.post_highway.parameters()) , lr=config['lr_gan_g']) #self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=config['lr_gan_d']) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss() def init_weights(self, m): if isinstance(m, nn.Linear): m.weight.data.uniform_(-0.02, 0.02) m.bias.data.fill_(0) def sample_post(self, x, c): xc = torch.cat((x, c), 1) e, mu, log_s = self.post_net(xc) #h_post = self.post_highway(xc) z, det_f, _, _ = self.post_generator((e, torch.eye(e.shape[1]), c, mu)) #h_prior = self.prior_highway(c) tilde_z, det_g, _ = self.prior_generator((z, det_f, c)) return tilde_z, z, mu, log_s, det_f, det_g def sample_code_post(self, x, c): xc = torch.cat((x, c), 1) e, mu, log_s = self.post_net(xc) #h_post = self.post_highway(xc) z, det_f, _, _ = self.post_generator((e, torch.eye(e.shape[1]), c, mu)) #h_prior = self.prior_highway(c) tilde_z, det_g, _ = self.prior_generator((z, det_f, c)) return tilde_z, mu, log_s, det_f, det_g def sample_post2(self, x, c): xc = torch.cat((x, c), 1) e, mu, log_s = self.post_net(xc) #h_post = self.post_highway(xc) z, det_f, _, _ = self.post_generator((e, torch.eye(e.shape[1]), c, mu)) return e, mu, log_s, z, det_f def sample_code_prior(self, c): e, mu, log_s = self.prior_net(c) #z = self.prior_generator(e) #h_prior = self.prior_highway(c) #tilde_z, det_g, _ = self.prior_generator((e, 0, h_prior)) return e, mu, log_s #, det_g def sample_prior(self, c): e, mu, log_s = self.prior_net(c) #h_prior = self.prior_highway(c) z, det_prior, _ = self.prior_generator((e, 0, c)) return z, det_prior def train_AE(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.train() self.decoder.train() c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) z, _, _, _, _ = self.sample_code_post(x, c) z_post, mu_post, log_s_post, det_f, det_g = self.sample_code_post(x, c) #prior_z, mu_prior, log_s_prior = self.sample_code_prior(c) #KL_loss = torch.sum(log_s_prior - log_s_post + (torch.exp(log_s_post) + (mu_post - mu_prior)**2)/torch.exp(log_s_prior),1) / 2 - 100 #kloss = KL_loss - det_f #+ det_g #KL_loss = log_Normal_diag(z_post, mu_post, log_s_post) - log_Normal_diag(prior_z, mu_prior, log_s_prior) output = self.decoder(torch.cat((z_post, c), 1), None, response[:, :-1], (res_lens - 1)) flattened_output = output.view(-1, self.vocab_size) dec_target = response[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # [(batch_sz*seq_len)] masked_target = dec_target.masked_select(mask) # output_mask = mask.unsqueeze(1).expand( mask.size(0), self.vocab_size) # [(batch_sz*seq_len) x n_tokens] masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) #print(KL_loss.mean()) #print(det_f.mean()) self.optimizer_AE.zero_grad() AE_term = self.criterion_ce(masked_output / self.temp, masked_target) loss = AE_term #+ KL_loss.mean() loss.backward() #torch.nn.utils.clip_grad_norm_(list(self.context_encoder.parameters())+list(self.decoder.parameters()), self.clip) torch.nn.utils.clip_grad_norm_( list(self.context_encoder.parameters()) + list(self.decoder.parameters()) + list(self.post_generator.parameters()) + list(self.prior_generator.parameters()) + list(self.post_net.parameters()), self.clip) self.optimizer_AE.step() return [ ('train_loss_AE', AE_term.item()) ] #,('KL_loss', KL_loss.mean().item())]#,('det_f', det_f.mean().item()),('det_g', det_g.mean().item())] def train_G(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.eval() self.optimizer_G.zero_grad() c = self.context_encoder(context, context_lens, utt_lens, floors) # -----------------posterior samples --------------------------- x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) z_0, mu_post, log_s_post, z_post, weight = self.sample_post2( x.detach(), c.detach()) # ----------------- prior samples --------------------------- prior_z, mu_prior, log_s_prior = self.sample_code_prior(c.detach()) KL_loss = torch.sum( log_s_prior - log_s_post + torch.exp(log_s_post) / torch.exp(log_s_prior) * torch.sum(weight**2, dim=2) + (mu_post)**2 / torch.exp(log_s_prior), 1) / 2 - 100 #KL_loss = abs(log_Normal_diag(z_0, mu_post, log_s_post) - log_Normal_diag(z_post, mu_prior, log_s_prior)) #KL_loss2 = torch.sum((prior_z - mu_post.detach())**2 / (2 * torch.exp(log_s_post.detach())),1) #print(mu_post.shape, prior_z.shape) loss = KL_loss #print(-det_f , KL_loss ) #loss = abs(loss) loss.mean().backward() torch.nn.utils.clip_grad_norm_( list(self.post_generator.parameters()) + list(self.prior_generator.parameters()) + list(self.post_net.parameters()) + list(self.prior_generator.parameters()), self.clip) self.optimizer_G.step() #costG = errG_prior - errG_post return [ ('KL_loss', KL_loss.mean().item()) ] #,('det_f', det_f.mean().item()),('det_g', det_g.sum().item())] def valid(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.eval() #self.discriminator.eval() self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) post_z, mu_post, log_s_post, det_f, det_g = self.sample_code_post(x, c) prior_z, mu_prior, log_s_prior = self.sample_code_prior(c) #errD_post = torch.mean(self.discriminator(torch.cat((post_z, c),1))) #errD_prior = torch.mean(self.discriminator(torch.cat((prior_z, c),1))) KL_loss = torch.sum( log_s_prior - log_s_post + (torch.exp(log_s_post) + (mu_post)**2) / torch.exp(log_s_prior), 1) / 2 #KL_loss = log_Normal_diag(post_z, mu_post, log_s_post) - log_Normal_diag(prior_z, mu_prior, log_s_prior) #KL_loss2 = torch.sum((prior_z - mu_post)**2 / (2 * torch.exp(log_s_post)),1) loss = KL_loss # -det_f costG = loss.sum() dec_target = response[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # [(batch_sz*seq_len)] masked_target = dec_target.masked_select(mask) output_mask = mask.unsqueeze(1).expand(mask.size(0), self.vocab_size) output = self.decoder(torch.cat((post_z, c), 1), None, response[:, :-1], (res_lens - 1)) flattened_output = output.view(-1, self.vocab_size) masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) lossAE = self.criterion_ce(masked_output / self.temp, masked_target) return [('valid_loss_AE', lossAE.item()), ('valid_loss_G', costG.item())] def sample(self, context, context_lens, utt_lens, floors, repeat, SOS_tok, EOS_tok): self.context_encoder.eval() self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) c_repeated = c.expand(repeat, -1) prior_z, _ = self.sample_prior(c_repeated) sample_words, sample_lens = self.decoder.sampling( torch.cat((prior_z, c_repeated), 1), None, self.maxlen, SOS_tok, EOS_tok, "greedy") return sample_words, sample_lens def gen(self, context, prior_z, context_lens, utt_lens, floors, repeat, SOS_tok, EOS_tok): self.context_encoder.eval() self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) c_repeated = c.expand(repeat, -1) sample_words, sample_lens = self.decoder.sampling( torch.cat((prior_z, c_repeated), 1), None, self.maxlen, SOS_tok, EOS_tok, "greedy") return sample_words, sample_lens def sample_latent(self, context, context_lens, utt_lens, floors, repeat, SOS_tok, EOS_tok): self.context_encoder.eval() #self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) c_repeated = c.expand(repeat, -1) e, _, _ = self.sample_code_prior(c_repeated) prior_z, _, _ = self.prior_generator((e, 0, c_repeated)) return prior_z, e def sample_latent_post(self, context, context_lens, utt_lens, floors, response, res_lens, repeat): self.context_encoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) c_repeated = c.expand(repeat, -1) x_repeated = x.expand(repeat, -1) z_post, z, mu_post, log_s_post, det_f, det_g = self.sample_post( x_repeated, c_repeated) return z_post, z def adjust_lr(self): self.lr_scheduler_AE.step()
def __init__(self, config, vocab_size, PAD_token=0): super(DFVAE, self).__init__() self.vocab_size = vocab_size self.maxlen = config['maxlen'] self.clip = config['clip'] self.lambda_gp = config['lambda_gp'] self.temp = config['temp'] self.embedder = nn.Embedding(vocab_size, config['emb_size'], padding_idx=PAD_token) self.utt_encoder = Encoder(self.embedder, config['emb_size'], config['n_hidden'], True, config['n_layers'], config['noise_radius']) self.context_encoder = ContextEncoder(self.utt_encoder, config['n_hidden'] * 2 + 2, config['n_hidden'], 1, config['noise_radius']) self.prior_net = Variation(config['n_hidden'], config['z_size']) # p(e|c) self.post_net = Variation(config['n_hidden'] * 3, config['z_size']) # q(e|c,x) #self.prior_highway = nn.Linear(config['n_hidden'], config['n_hidden']) #self.post_highway = nn.Linear(config['n_hidden'] * 3, config['n_hidden']) self.postflow1 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.postflow2 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.postflow3 = flow.myIAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow1 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow2 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.priorflow3 = flow.IAF(config['z_size'], config['z_size'] * 2, config['n_hidden'], 3) self.post_generator = nn_.SequentialFlow(self.postflow1, self.postflow2, self.postflow3) self.prior_generator = nn_.SequentialFlow(self.priorflow1, self.priorflow2, self.priorflow3) self.decoder = Decoder(self.embedder, config['emb_size'], config['n_hidden'] + config['z_size'], vocab_size, n_layers=1) self.optimizer_AE = optim.SGD( list(self.context_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.decoder.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()) #+list(self.prior_highway.parameters()) #+list(self.post_highway.parameters()) , lr=config['lr_ae']) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()) #+list(self.prior_highway.parameters()) #+list(self.post_highway.parameters()) , lr=config['lr_gan_g']) #self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=config['lr_gan_d']) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss()
def __init__(self, config, vocab_size, PAD_token=0): super(DialogWAE, self).__init__() self.vocab_size = vocab_size self.maxlen = config['maxlen'] self.clip = config['clip'] self.lambda_gp = config['lambda_gp'] self.temp = config['temp'] self.embedder = nn.Embedding(vocab_size, config['emb_size'], padding_idx=PAD_token) self.utt_encoder = Encoder(self.embedder, config['emb_size'], config['n_hidden'], True, config['n_layers'], config['noise_radius']) self.context_encoder = ContextEncoder(self.utt_encoder, config['n_hidden'] * 2 + 2, config['n_hidden'], 1, config['noise_radius']) self.prior_net = Variation(config['n_hidden'], config['z_size']) # p(e|c) self.post_net = Variation(config['n_hidden'] * 3, config['z_size']) # q(e|c,x) self.post_generator = nn.Sequential( nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size'])) self.post_generator.apply(self.init_weights) self.prior_generator = nn.Sequential( nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size'])) self.prior_generator.apply(self.init_weights) self.decoder = Decoder(self.embedder, config['emb_size'], config['n_hidden'] + config['z_size'], vocab_size, n_layers=1) self.discriminator = nn.Sequential( nn.Linear(config['n_hidden'] + config['z_size'], config['n_hidden'] * 2), nn.BatchNorm1d(config['n_hidden'] * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config['n_hidden'] * 2, config['n_hidden'] * 2), nn.BatchNorm1d(config['n_hidden'] * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config['n_hidden'] * 2, 1), ) self.discriminator.apply(self.init_weights) self.optimizer_AE = optim.SGD(list(self.context_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.decoder.parameters()), lr=config['lr_ae']) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()), lr=config['lr_gan_g']) self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=config['lr_gan_d']) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss()
class DialogWAE(nn.Module): def __init__(self, config, vocab_size, PAD_token=0): super(DialogWAE, self).__init__() self.vocab_size = vocab_size self.maxlen = config['maxlen'] self.clip = config['clip'] self.lambda_gp = config['lambda_gp'] self.temp = config['temp'] self.embedder = nn.Embedding(vocab_size, config['emb_size'], padding_idx=PAD_token) self.utt_encoder = Encoder(self.embedder, config['emb_size'], config['n_hidden'], True, config['n_layers'], config['noise_radius']) self.context_encoder = ContextEncoder(self.utt_encoder, config['n_hidden'] * 2 + 2, config['n_hidden'], 1, config['noise_radius']) self.prior_net = Variation(config['n_hidden'], config['z_size']) # p(e|c) self.post_net = Variation(config['n_hidden'] * 3, config['z_size']) # q(e|c,x) self.post_generator = nn.Sequential( nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size'])) self.post_generator.apply(self.init_weights) self.prior_generator = nn.Sequential( nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size']), nn.BatchNorm1d(config['z_size'], eps=1e-05, momentum=0.1), nn.ReLU(), nn.Linear(config['z_size'], config['z_size'])) self.prior_generator.apply(self.init_weights) self.decoder = Decoder(self.embedder, config['emb_size'], config['n_hidden'] + config['z_size'], vocab_size, n_layers=1) self.discriminator = nn.Sequential( nn.Linear(config['n_hidden'] + config['z_size'], config['n_hidden'] * 2), nn.BatchNorm1d(config['n_hidden'] * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config['n_hidden'] * 2, config['n_hidden'] * 2), nn.BatchNorm1d(config['n_hidden'] * 2, eps=1e-05, momentum=0.1), nn.LeakyReLU(0.2), nn.Linear(config['n_hidden'] * 2, 1), ) self.discriminator.apply(self.init_weights) self.optimizer_AE = optim.SGD(list(self.context_encoder.parameters()) + list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.decoder.parameters()), lr=config['lr_ae']) self.optimizer_G = optim.RMSprop( list(self.post_net.parameters()) + list(self.post_generator.parameters()) + list(self.prior_net.parameters()) + list(self.prior_generator.parameters()), lr=config['lr_gan_g']) self.optimizer_D = optim.RMSprop(self.discriminator.parameters(), lr=config['lr_gan_d']) self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss() def init_weights(self, m): if isinstance(m, nn.Linear): m.weight.data.uniform_(-0.02, 0.02) m.bias.data.fill_(0) def sample_code_post(self, x, c): e, _, _ = self.post_net(torch.cat((x, c), 1)) z = self.post_generator(e) return z def sample_code_prior(self, c): e, _, _ = self.prior_net(c) z = self.prior_generator(e) return z def train_AE(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.train() self.decoder.train() c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) z = self.sample_code_post(x, c) output = self.decoder(torch.cat((z, c), 1), None, response[:, :-1], (res_lens - 1)) flattened_output = output.view(-1, self.vocab_size) dec_target = response[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # [(batch_sz*seq_len)] masked_target = dec_target.masked_select(mask) # output_mask = mask.unsqueeze(1).expand( mask.size(0), self.vocab_size) # [(batch_sz*seq_len) x n_tokens] masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) self.optimizer_AE.zero_grad() loss = self.criterion_ce(masked_output / self.temp, masked_target) loss.backward() torch.nn.utils.clip_grad_norm_( list(self.context_encoder.parameters()) + list(self.decoder.parameters()), self.clip) self.optimizer_AE.step() return [('train_loss_AE', loss.item())] def train_G(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.eval() self.optimizer_G.zero_grad() for p in self.discriminator.parameters(): p.requires_grad = False c = self.context_encoder(context, context_lens, utt_lens, floors) # -----------------posterior samples --------------------------- x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) z_post = self.sample_code_post(x.detach(), c.detach()) errG_post = torch.mean( self.discriminator(torch.cat((z_post, c.detach()), 1))) errG_post.backward(minus_one) # ----------------- prior samples --------------------------- prior_z = self.sample_code_prior(c.detach()) errG_prior = torch.mean( self.discriminator(torch.cat((prior_z, c.detach()), 1))) errG_prior.backward(one) self.optimizer_G.step() for p in self.discriminator.parameters(): p.requires_grad = True costG = errG_prior - errG_post return [('train_loss_G', costG.item())] def train_D(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.eval() self.discriminator.train() self.optimizer_D.zero_grad() batch_size = context.size(0) c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) post_z = self.sample_code_post(x, c) errD_post = torch.mean( self.discriminator(torch.cat((post_z.detach(), c.detach()), 1))) errD_post.backward(one) prior_z = self.sample_code_prior(c) errD_prior = torch.mean( self.discriminator(torch.cat((prior_z.detach(), c.detach()), 1))) errD_prior.backward(minus_one) alpha = gData(torch.rand(batch_size, 1)) alpha = alpha.expand(prior_z.size()) interpolates = alpha * prior_z.data + ((1 - alpha) * post_z.data) interpolates = Variable(interpolates, requires_grad=True) d_input = torch.cat((interpolates, c.detach()), 1) disc_interpolates = torch.mean(self.discriminator(d_input)) gradients = torch.autograd.grad( outputs=disc_interpolates, inputs=interpolates, grad_outputs=gData(torch.ones(disc_interpolates.size())), create_graph=True, retain_graph=True, only_inputs=True)[0] gradient_penalty = ( (gradients.contiguous().view(gradients.size(0), -1).norm(2, dim=1) - 1)**2).mean() * self.lambda_gp gradient_penalty.backward() self.optimizer_D.step() costD = -(errD_prior - errD_post) + gradient_penalty return [('train_loss_D', costD.item())] def valid(self, context, context_lens, utt_lens, floors, response, res_lens): self.context_encoder.eval() self.discriminator.eval() self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) x, _ = self.utt_encoder(response[:, 1:], res_lens - 1) post_z = self.sample_code_post(x, c) prior_z = self.sample_code_prior(c) errD_post = torch.mean(self.discriminator(torch.cat((post_z, c), 1))) errD_prior = torch.mean(self.discriminator(torch.cat((prior_z, c), 1))) costD = -(errD_prior - errD_post) costG = -costD dec_target = response[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # [(batch_sz*seq_len)] masked_target = dec_target.masked_select(mask) output_mask = mask.unsqueeze(1).expand(mask.size(0), self.vocab_size) output = self.decoder(torch.cat((post_z, c), 1), None, response[:, :-1], (res_lens - 1)) flattened_output = output.view(-1, self.vocab_size) masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) lossAE = self.criterion_ce(masked_output / self.temp, masked_target) return [('valid_loss_AE', lossAE.item()), ('valid_loss_G', costG.item()), ('valid_loss_D', costD.item())] def sample(self, context, context_lens, utt_lens, floors, repeat, SOS_tok, EOS_tok): self.context_encoder.eval() self.decoder.eval() c = self.context_encoder(context, context_lens, utt_lens, floors) # encode context into embedding c_repeated = c.expand(repeat, -1) prior_z = self.sample_code_prior(c_repeated) # print(prior_z.shape) # print(prior_z) sample_words, sample_lens = self.decoder.sampling( torch.cat((prior_z, c_repeated), 1), None, self.maxlen, SOS_tok, EOS_tok, "greedy") return sample_words, sample_lens def adjust_lr(self): self.lr_scheduler_AE.step()
def __init__(self, config, api, PAD_token=0): super(CVAE, self).__init__() self.vocab = api.vocab self.vocab_size = len(self.vocab) self.rev_vocab = api.rev_vocab self.go_id = self.rev_vocab["<s>"] self.eos_id = self.rev_vocab["</s>"] self.maxlen = config.maxlen self.clip = config.clip self.temp = config.temp self.full_kl_step = config.full_kl_step self.z_size = config.z_size self.init_w = config.init_weight self.softmax = nn.Softmax(dim=1) self.embedder = nn.Embedding(self.vocab_size, config.emb_size, padding_idx=PAD_token) # 对title, 每一句诗做编码 self.seq_encoder = Encoder(embedder=self.embedder, input_size=config.emb_size, hidden_size=config.n_hidden, bidirectional=True, n_layers=config.n_layers, noise_radius=config.noise_radius) # 先验网络的输入是 标题encode结果 + 上一句诗过encoder的结果 + 上一句情感过encoder的结果 self.prior_net = Variation(config.n_hidden * 4, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # 后验网络,再加上x的2*hidden # self.post_net = Variation(config.n_hidden * 6, config.z_size*2) self.post_net = Variation(config.n_hidden * 6, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # 词包loss的MLP self.bow_project = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, 400), nn.LeakyReLU(), nn.Dropout(config.dropout), nn.Linear(400, self.vocab_size)) self.init_decoder_hidden = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden), nn.BatchNorm1d(config.n_hidden, eps=1e-05, momentum=0.1), nn.LeakyReLU()) # self.post_generator = nn.Sequential( # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(config.z_size, config.z_size) # ) # self.post_generator.apply(self.init_weights) # self.prior_generator = nn.Sequential( # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(config.dropout), # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(config.dropout), # nn.Linear(config.z_size, config.z_size) # ) # self.prior_generator.apply(self.init_weights) self.init_decoder_hidden.apply(self.init_weights) self.bow_project.apply(self.init_weights) self.post_net.apply(self.init_weights) self.decoder = Decoder(embedder=self.embedder, input_size=config.emb_size, hidden_size=config.n_hidden, vocab_size=self.vocab_size, n_layers=1) # self.optimizer_lead = optim.Adam(list(self.seq_encoder.parameters())\ # + list(self.prior_net.parameters()), lr=config.lr_lead) self.optimizer_AE = optim.Adam(list(self.seq_encoder.parameters())\ + list(self.prior_net.parameters())\ # + list(self.prior_generator.parameters()) + list(self.post_net.parameters())\ # + list(self.post_generator.parameters()) + list(self.bow_project.parameters())\ + list(self.init_decoder_hidden.parameters())\ + list(self.decoder.parameters()), lr=config.lr_ae) # self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss() self.softmax = nn.Softmax(dim=1) self.criterion_sent_lead = nn.CrossEntropyLoss()
class CVAE(nn.Module): def __init__(self, config, api, PAD_token=0): super(CVAE, self).__init__() self.vocab = api.vocab self.vocab_size = len(self.vocab) self.rev_vocab = api.rev_vocab self.go_id = self.rev_vocab["<s>"] self.eos_id = self.rev_vocab["</s>"] self.maxlen = config.maxlen self.clip = config.clip self.temp = config.temp self.full_kl_step = config.full_kl_step self.z_size = config.z_size self.init_w = config.init_weight self.softmax = nn.Softmax(dim=1) self.embedder = nn.Embedding(self.vocab_size, config.emb_size, padding_idx=PAD_token) # 对title, 每一句诗做编码 self.seq_encoder = Encoder(embedder=self.embedder, input_size=config.emb_size, hidden_size=config.n_hidden, bidirectional=True, n_layers=config.n_layers, noise_radius=config.noise_radius) # 先验网络的输入是 标题encode结果 + 上一句诗过encoder的结果 + 上一句情感过encoder的结果 self.prior_net = Variation(config.n_hidden * 4, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # 后验网络,再加上x的2*hidden # self.post_net = Variation(config.n_hidden * 6, config.z_size*2) self.post_net = Variation(config.n_hidden * 6, config.z_size, dropout_rate=config.dropout, init_weight=self.init_w) # 词包loss的MLP self.bow_project = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, 400), nn.LeakyReLU(), nn.Dropout(config.dropout), nn.Linear(400, self.vocab_size)) self.init_decoder_hidden = nn.Sequential( nn.Linear(config.n_hidden * 4 + config.z_size, config.n_hidden), nn.BatchNorm1d(config.n_hidden, eps=1e-05, momentum=0.1), nn.LeakyReLU()) # self.post_generator = nn.Sequential( # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(config.z_size, config.z_size) # ) # self.post_generator.apply(self.init_weights) # self.prior_generator = nn.Sequential( # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(config.dropout), # nn.Linear(config.z_size, config.z_size), # nn.BatchNorm1d(config.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(config.dropout), # nn.Linear(config.z_size, config.z_size) # ) # self.prior_generator.apply(self.init_weights) self.init_decoder_hidden.apply(self.init_weights) self.bow_project.apply(self.init_weights) self.post_net.apply(self.init_weights) self.decoder = Decoder(embedder=self.embedder, input_size=config.emb_size, hidden_size=config.n_hidden, vocab_size=self.vocab_size, n_layers=1) # self.optimizer_lead = optim.Adam(list(self.seq_encoder.parameters())\ # + list(self.prior_net.parameters()), lr=config.lr_lead) self.optimizer_AE = optim.Adam(list(self.seq_encoder.parameters())\ + list(self.prior_net.parameters())\ # + list(self.prior_generator.parameters()) + list(self.post_net.parameters())\ # + list(self.post_generator.parameters()) + list(self.bow_project.parameters())\ + list(self.init_decoder_hidden.parameters())\ + list(self.decoder.parameters()), lr=config.lr_ae) # self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss() self.softmax = nn.Softmax(dim=1) self.criterion_sent_lead = nn.CrossEntropyLoss() def set_full_kl_step(self, kl_full_step): self.full_kl_step = kl_full_step def force_change_lr(self, new_init_lr_ae): self.optimizer_AE = optim.Adam(list(self.seq_encoder.parameters()) \ + list(self.prior_net.parameters()) \ # + list(self.prior_generator.parameters()) + list(self.post_net.parameters()) \ # + list(self.post_generator.parameters()) + list(self.bow_project.parameters()) \ + list(self.init_decoder_hidden.parameters()) \ + list(self.decoder.parameters()), lr=new_init_lr_ae) def init_weights(self, m): if isinstance(m, nn.Linear): m.weight.data.uniform_(-self.init_w, self.init_w) m.bias.data.fill_(0) def sample_code_post(self, x, c): # import pdb # pdb.set_trace() # mulogsigma = self.post_net(torch.cat((x, c), dim=1)) # mu, logsigma = torch.chunk(mulogsigma, chunks=2, dim=1) # batch_size = c.size(0) # std = torch.exp(0.5 * logsigma) # epsilon = to_tensor(torch.randn([batch_size, self.z_size])) # z = epsilon * std + mu z, mu, logsigma = self.post_net(torch.cat( (x, c), 1)) # 输入:(batch, 3*2*n_hidden) # z = self.post_generator(z) return z, mu, logsigma def sample_code_prior(self, c, sentiment_mask=None, mask_type=None): return self.prior_net(c, sentiment_mask=sentiment_mask, mask_type=mask_type) # # input: (batch, 3) # # target: (batch, 3) # def criterion_sent_lead(self, input, target): # softmax_res = self.softmax(input) # negative_log_softmax_res = -torch.log(softmax_res + 1e-10) # (batch, 3) # cross_entropy_loss = torch.sum(negative_log_softmax_res * target, dim=1) # avg_cross_entropy = torch.mean(cross_entropy_loss) # return avg_cross_entropy # sentiment_lead: (batch, 3) def train_AE(self, global_t, title, context, target, target_lens, sentiment_mask=None, sentiment_lead=None): self.seq_encoder.train() self.decoder.train() # batch_size = title.size(0) # 每一句的情感用第二个分类器来预测,输入当前的m_hidden,输出分类结果 title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) # import pdb # pdb.set_trace() x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) condition_prior = torch.cat((title_last_hidden, context_last_hidden), dim=1) z_prior, prior_mu, prior_logvar, pi, pi_final = self.sample_code_prior( condition_prior, sentiment_mask=sentiment_mask) z_post, post_mu, post_logvar = self.sample_code_post( x, condition_prior) # import pdb # pdb.set_trace() if sentiment_lead is not None: self.sent_lead_loss = self.criterion_sent_lead( input=pi, target=sentiment_lead) else: self.sent_lead_loss = 0 # if sentiment_lead is not None: # self.optimizer_lead.zero_grad() # self.sent_lead_loss.backward() # self.optimizer_lead.step() # return [('lead_loss', self.sent_lead_loss.item())], global_t final_info = torch.cat((z_post, condition_prior), dim=1) # reconstruct_loss # import pdb # pdb.set_trace() output = self.decoder(init_hidden=self.init_decoder_hidden(final_info), context=None, inputs=target[:, :-1]) flattened_output = output.view(-1, self.vocab_size) # flattened_output = self.softmax(flattened_output) + 1e-10 # flattened_output = torch.log(flattened_output) dec_target = target[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # 即判断target的token中是否有0(pad项) masked_target = dec_target.masked_select(mask) # 选出非pad项 output_mask = mask.unsqueeze(1).expand( mask.size(0), self.vocab_size) # [(batch_sz * seq_len) x n_tokens] masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) self.rc_loss = self.criterion_ce(masked_output / self.temp, masked_target) # kl散度 kld = gaussian_kld(post_mu, post_logvar, prior_mu, prior_logvar) self.avg_kld = torch.mean(kld) self.kl_weights = min(global_t / self.full_kl_step, 1.0) # 退火 self.kl_loss = self.kl_weights * self.avg_kld # avg_bow_loss self.bow_logits = self.bow_project(final_info) # 说白了就是把target所有词的预测loss求个和 labels = target[:, 1:] label_mask = torch.sign(labels).detach().float() # 取符号变成正数,从而通过最小化来optimize # soft_result = self.softmax(self.bow_logits) + 1e-10 # bow_loss = -torch.log(soft_result).gather(1, labels) * label_mask bow_loss = -F.log_softmax(self.bow_logits, dim=1).gather( 1, labels) * label_mask sum_bow_loss = torch.sum(bow_loss, 1) self.avg_bow_loss = torch.mean(sum_bow_loss) self.aug_elbo_loss = self.avg_bow_loss + self.kl_loss + self.rc_loss self.total_loss = self.aug_elbo_loss + self.sent_lead_loss # 变相增加标注集的学习率 if sentiment_mask is not None: self.total_loss = self.total_loss * 13.33 self.optimizer_AE.zero_grad() self.total_loss.backward() self.optimizer_AE.step() avg_total_loss = self.total_loss.item() avg_lead_loss = 0 if sentiment_lead is None else self.sent_lead_loss.item( ) avg_aug_elbo_loss = self.aug_elbo_loss.item() avg_kl_loss = self.kl_loss.item() avg_rc_loss = self.rc_loss.data.item() avg_bow_loss = self.avg_bow_loss.item() global_t += 1 return [('avg_total_loss', avg_total_loss), ('avg_lead_loss', avg_lead_loss), ('avg_aug_elbo_loss', avg_aug_elbo_loss), ('avg_kl_loss', avg_kl_loss), ('avg_rc_loss', avg_rc_loss), ('avg_bow_loss', avg_bow_loss), ('kl_weight', self.kl_weights)], global_t def valid_AE(self, global_t, title, context, target, target_lens, sentiment_mask=None, sentiment_lead=None): self.seq_encoder.eval() self.decoder.eval() title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) # import pdb # pdb.set_trace() x, _ = self.seq_encoder(target[:, 1:], target_lens - 1) condition_prior = torch.cat((title_last_hidden, context_last_hidden), dim=1) z_prior, prior_mu, prior_logvar, pi, pi_final = self.sample_code_prior( condition_prior, sentiment_mask=sentiment_mask) z_post, post_mu, post_logvar = self.sample_code_post( x, condition_prior) if sentiment_lead is not None: self.sent_lead_loss = self.criterion_sent_lead( input=pi, target=sentiment_lead) else: self.sent_lead_loss = 0 # if sentiment_lead is not None: # return [('valid_lead_loss', self.sent_lead_loss.item())], global_t final_info = torch.cat((z_post, condition_prior), dim=1) output = self.decoder(init_hidden=self.init_decoder_hidden(final_info), context=None, inputs=target[:, :-1]) flattened_output = output.view(-1, self.vocab_size) # flattened_output = self.softmax(flattened_output) + 1e-10 # flattened_output = torch.log(flattened_output) dec_target = target[:, 1:].contiguous().view(-1) mask = dec_target.gt(0) # 即判断target的token中是否有0(pad项) masked_target = dec_target.masked_select(mask) # 选出非pad项 output_mask = mask.unsqueeze(1).expand( mask.size(0), self.vocab_size) # [(batch_sz * seq_len) x n_tokens] masked_output = flattened_output.masked_select(output_mask).view( -1, self.vocab_size) self.rc_loss = self.criterion_ce(masked_output / self.temp, masked_target) # kl散度 kld = gaussian_kld(post_mu, post_logvar, prior_mu, prior_logvar) self.avg_kld = torch.mean(kld) self.kl_weights = min(global_t / self.full_kl_step, 1.0) # 退火 self.kl_loss = self.kl_weights * self.avg_kld # avg_bow_loss self.bow_logits = self.bow_project(final_info) # 说白了就是把target所有词的预测loss求个和 labels = target[:, 1:] label_mask = torch.sign(labels).detach().float() bow_loss = -F.log_softmax(self.bow_logits, dim=1).gather( 1, labels) * label_mask sum_bow_loss = torch.sum(bow_loss, 1) self.avg_bow_loss = torch.mean(sum_bow_loss) self.aug_elbo_loss = self.avg_bow_loss + self.kl_loss + self.rc_loss avg_aug_elbo_loss = self.aug_elbo_loss.item() avg_kl_loss = self.kl_loss.item() avg_rc_loss = self.rc_loss.data.item() avg_bow_loss = self.avg_bow_loss.item() avg_lead_loss = 0 if sentiment_lead is None else self.sent_lead_loss.item( ) return [('valid_lead_loss', avg_lead_loss), ('valid_aug_elbo_loss', avg_aug_elbo_loss), ('valid_kl_loss', avg_kl_loss), ('valid_rc_loss', avg_rc_loss), ('valid_bow_loss', avg_bow_loss)], global_t # batch_size = 1 只输入了一个标题 # test的时候,只有先验,没有后验,更没有所谓的kl散度 def test(self, title_tensor, title_words, mask_type=None): self.seq_encoder.eval() self.decoder.eval() assert title_tensor.size(0) == 1 tem = [[2, 3] + [0] * (self.maxlen - 2)] pred_poems = [] # 过滤掉标题中的<s> </s> 0,只为了打印 title_tokens = [ self.vocab[e] for e in title_words[0].tolist() if e not in [0, self.eos_id, self.go_id] ] pred_poems.append(title_tokens) for i in range(4): tem = to_tensor(np.array(tem)) context = tem if i == 0: context_last_hidden, _ = self.seq_encoder(title_tensor) else: context_last_hidden, _ = self.seq_encoder(context) title_last_hidden, _ = self.seq_encoder(title_tensor) condition_prior = torch.cat( (title_last_hidden, context_last_hidden), dim=1) z_prior, prior_mu, prior_logvar, _, _ = self.sample_code_prior( condition_prior, mask_type=mask_type) final_info = torch.cat((z_prior, condition_prior), 1) decode_words = self.decoder.testing( init_hidden=self.init_decoder_hidden(final_info), maxlen=self.maxlen, go_id=self.go_id, mode="greedy") decode_words = decode_words[0].tolist() if len(decode_words) >= self.maxlen: tem = [decode_words[0:self.maxlen]] else: tem = [[0] * (self.maxlen - len(decode_words)) + decode_words] pred_tokens = [ self.vocab[e] for e in decode_words[:-1] if e != self.eos_id and e != 0 and e != self.go_id ] pred_poems.append(pred_tokens) gen = "" for line in pred_poems: cur_line = " ".join(line) gen = gen + cur_line + '\n' return gen def sample(self, title, context, repeat, go_id, end_id): self.seq_encoder.eval() self.decoder.eval() assert title.size(0) == 1 title_last_hidden, _ = self.seq_encoder(title) context_last_hidden, _ = self.seq_encoder(context) condition_prior = torch.cat((title_last_hidden, context_last_hidden), 1) condition_prior_repeat = condition_prior.expand(repeat, -1) z_prior_repeat, _, _, _, _ = self.sample_code_prior( condition_prior_repeat) final_info = torch.cat((z_prior_repeat, condition_prior_repeat), dim=1) sample_words, sample_lens = self.decoder.sampling( init_hidden=self.init_decoder_hidden(final_info), maxlen=self.maxlen, go_id=self.go_id, eos_id=self.eos_id, mode="greedy") return sample_words, sample_lens
def __init__(self, config, api, PAD_token=0): super(CVAE, self).__init__() assert api.rev_vocab['<pad>'] == PAD_token self.vocab = api.vocab self.vocab_size = len(self.vocab) self.embed_size = config.emb_size self.sent_class_size = config.sent_class self.sent_emb_size = config.sent_emb_size self.hidden_size = config.n_hidden self.bow_size = config.bow_size self.rev_vocab = api.rev_vocab self.dropout = config.dropout self.go_id = self.rev_vocab["<s>"] self.eos_id = self.rev_vocab["</s>"] self.maxlen = config.maxlen self.clip = config.clip self.temp = config.temp self.full_kl_step = config.full_kl_step self.z_size = config.z_size self.init_w = config.init_weight self.softmax = nn.Softmax(dim=1) self.bidirectional = config.bidirectional self.lr_ae = config.lr_ae # 如果LSTM双向,则两个方向拼接在一起 self.encoder_output_size = self.hidden_size * (1 + int(self.bidirectional)) # 标题和首句拼接在一起,得到cvae的condition部分 self.prior_input_dim = self.encoder_output_size * 2 # 在prior的基础上再拼接对target的encode结果 self.post_input_dim = self.prior_input_dim + self.encoder_output_size self.decoder_input_size = self.prior_input_dim + self.z_size + self.sent_emb_size self.embedder = nn.Embedding(self.vocab_size, self.embed_size, padding_idx=PAD_token) self.sent_embedder = nn.Embedding(self.sent_class_size, self.sent_emb_size) # 对title, 每一句诗做编码, 默认双向LSTM,将最终的一维拼在一起 self.seq_encoder = Encoder(embedder=self.embedder, input_size=config.emb_size, hidden_size=config.n_hidden, bidirectional=self.bidirectional, n_layers=config.n_layers, noise_radius=config.noise_radius) # 先验网络 self.prior_net = Variation(self.sent_class_size, self.sent_emb_size, self.prior_input_dim, self.z_size, dropout_rate=self.dropout, init_weight=self.init_w) # 后验网络 self.post_net = Variation(self.sent_class_size, self.sent_emb_size, self.post_input_dim, self.z_size, dropout_rate=self.dropout, init_weight=self.init_w) # 词包loss的MLP self.bow_project = nn.Sequential( nn.Linear(self.decoder_input_size, self.bow_size), nn.LeakyReLU(), nn.Dropout(self.dropout), nn.Linear(self.bow_size, self.vocab_size)) self.decoder = Decoder(embedder=self.embedder, input_size=self.embed_size, hidden_size=self.hidden_size, vocab_size=self.vocab_size, n_layers=1) self.init_decoder_hidden = nn.Sequential( nn.Linear(self.decoder_input_size, self.hidden_size), nn.BatchNorm1d(self.hidden_size, eps=1e-05, momentum=0.1), nn.LeakyReLU()) # self.post_generator = nn.Sequential( # nn.Linear(self.z_size, self.z_size), # nn.BatchNorm1d(self.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(self.z_size, self.z_size), # nn.BatchNorm1d(self.z_size, eps=1e-05, momentum=0.1), # nn.LeakyReLU(), # nn.Linear(self.z_size, self.z_size) # ) # self.post_generator.apply(self.init_weights) # self.prior_generator = nn.Sequential( # nn.Linear(self.z_size, self.z_size), # nn.BatchNorm1d(self.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(self.dropout), # nn.Linear(self.z_size, self.z_size), # nn.BatchNorm1d(self.z_size, eps=1e-05, momentum=0.1), # nn.ReLU(), # nn.Dropout(self.dropout), # nn.Linear(self.z_size, self.z_size) # ) # self.prior_generator.apply(self.init_weights) self.init_decoder_hidden.apply(self.init_weights) self.bow_project.apply(self.init_weights) self.post_net.apply(self.init_weights) # self.optimizer_lead = optim.Adam(list(self.seq_encoder.parameters())\ # + list(self.prior_net.parameters()), lr=self.lr_lead) self.optimizer_AE = optim.AdamW(self.parameters(), lr=self.lr_ae) # self.lr_scheduler_AE = optim.lr_scheduler.StepLR(self.optimizer_AE, step_size=10, gamma=0.6) self.criterion_ce = nn.CrossEntropyLoss() self.softmax = nn.Softmax(dim=1) self.criterion_sent_lead = nn.CrossEntropyLoss()