def forward(self, x):
"""
Params:
-------
sentence: sequence of word indices.
Returns:
--------
recon_loss: reconstruction loss of VAE.
kl_loss: KL-div loss of VAE.
"""
# Encode
x = to_gpu(x)
mu, logvar = self.forward_encoder(x)
# sample and decode
z = self.sample_z(mu, logvar, n_times=mu.size(0))
# z = self.sample_z(mu, logvar) # single eps for whole batch
y = self.forward_decoder(x, z)
# get losses
recon_loss = self.get_recon_loss(x, y)
kl_loss = self.get_kl_loss(mu, logvar)
return {
'recon_loss': recon_loss,
'kl_loss': kl_loss,
'mu': mu,
'logvar': logvar,
'z': z,
'y': y
}
def forward_multiple(self, x, n_times=16):
# encode
x = to_gpu(x)
mu, logvar = self.forward_encoder(x)
# sample and decode
zs = []
ys = []
kl_losses = []
recon_losses = []
for i in range(mu.size(0)):
z = self.sample_z(mu[i], logvar[i], n_times=n_times)
target = x[:,i:i+1].repeat(1, n_times)
y = self.forward_decoder(target, z)
recon_losses.append(self.get_recon_loss(target, y))
kl_losses.append(self.get_kl_loss(mu[i:i+1], logvar[i:i+1]))
return {
'recon_losses': recon_losses,
'kl_losses': kl_losses,
'mu': mu,
'logvar': logvar,
'zs': zs,
'ys': ys
}
def sample_z(self, mu, logvar, n_times=1):
"""
Reparameterization trick: z = mu + std*eps; eps ~ N(0, I)
"""
eps = Variable(torch.randn(n_times, self.z_dim))
eps = to_gpu(eps)
return mu + torch.exp(logvar/2) * eps
def get_recon_loss(self, sentence, y):
mbsize = sentence.size(1)
pad_words = Variable(torch.LongTensor([self.PAD_IDX])).repeat(1, mbsize)
pad_words = to_gpu(pad_words)
dec_targets = torch.cat([sentence[1:], pad_words], dim=0)
recon_loss = F.cross_entropy(
y.view(-1, self.n_vocab), dec_targets.view(-1), size_average=True, weight=self.weights
)
return recon_loss
def forward_multiple(self, x, n_times=16):
x = to_gpu(x)
ys = []
recon_losses = []
for i in range(x.size(1)):
target = x[:, i:i + 1].repeat(1, n_times)
y = self.forward(target)['y']
recon_losses.append(self.get_recon_loss(target, y))
ys.append(y)
return {'recon_losses': recon_losses, 'ys': ys}
def forward(self, x):
x = to_gpu(x)
inputs_emb = self.word_emb(x)
outputs, _ = self.lm(inputs_emb, None)
seq_len, mbsize, _ = outputs.size()
outputs = outputs.view(seq_len * mbsize, -1)
y = self.decode(outputs)
y = y.view(seq_len, mbsize, self.n_vocab)
recon_loss = self.get_recon_loss(x, y)
return {'recon_loss': recon_loss, 'y': y}
def sample_sentence(self, z, raw=False, temp=1):
"""
Sample single sentence from p(x|z,c) according to given temperature.
`raw = True` means this returns sentence as in dataset which is useful
to train discriminator. `False` means that this will return list of
`word_idx` which is useful for evaluation.
"""
word = torch.LongTensor([self.START_IDX])
word = word.cuda() if self.gpu else word
word = Variable(word) # '<start>'
z = z.view(1, 1, -1)
h = z
if not isinstance(h, Variable):
h = Variable(h)
outputs = []
if raw:
outputs.append(self.START_IDX)
for i in range(self.MAX_SENT_LEN):
emb = self.word_emb(word).view(1, 1, -1)
emb = torch.cat([emb, z], 2)
output, h = self.decoder(emb, h)
y = self.decode(output.view((1,-1))).view(-1)
y = F.softmax(y/temp, dim=0)
idx = torch.multinomial(y, num_samples=1)
word = Variable(torch.LongTensor([int(idx)]))
word = word.cuda() if self.gpu else word
idx = int(idx)
if not raw and idx == self.EOS_IDX:
break
outputs.append(idx)
# Back to default state: train
if raw:
outputs = Variable(torch.LongTensor(outputs)).unsqueeze(0)
return to_gpu(outputs)
else:
return outputs
def __init__(self,
n_vocab,
h_dim=128,
unk_idx=0,
pad_idx=1,
start_idx=2,
eos_idx=3,
max_sent_len=15,
use_input_embeddings=True,
set_other_to_random=False,
set_unk_to_random=True,
decode_with_embeddings=True,
rnn_dropout=0.3,
mask_pad=True,
pretrained_embeddings=None,
freeze_embeddings=False,
gpu=False):
super(RNN_LanguageModel, self).__init__()
self.PAD_IDX = pad_idx
self.UNK_IDX = unk_idx
self.START_IDX = start_idx
self.EOS_IDX = eos_idx
self.MAX_SENT_LEN = max_sent_len
self.freeze_embeddings = freeze_embeddings
self.n_vocab = n_vocab
self.h_dim = h_dim
self.use_input_embeddings = use_input_embeddings # in case we go for sentence embeddings
self.decode_with_embeddings = decode_with_embeddings
self.gpu = gpu
if mask_pad:
self.weights = np.ones(n_vocab)
self.weights[self.PAD_IDX] = 0
self.weights = to_gpu(torch.tensor(self.weights.astype(
np.float32)))
else:
self.weights = None
"""
Word embeddings layer
"""
if self.use_input_embeddings:
if pretrained_embeddings is None:
self.emb_dim = h_dim
self.word_emb = nn.Embedding(n_vocab, h_dim, self.PAD_IDX)
else:
self.emb_dim = pretrained_embeddings.size(1)
self.word_emb = nn.Embedding(n_vocab, self.emb_dim,
self.PAD_IDX)
# Set pretrained embeddings
self.word_emb.weight.data.copy_(pretrained_embeddings)
else:
raise NotImplementedError('not supported yet')
if set_unk_to_random:
self.word_emb.weight[self.UNK_IDX].data.copy_(
torch.tensor(np.random.randn(self.emb_dim)))
if set_other_to_random:
self.word_emb.weight[self.PAD_IDX].data.copy_(
torch.tensor(np.random.randn(self.emb_dim)))
self.word_emb.weight[self.START_IDX].data.copy_(
torch.tensor(np.random.randn(self.emb_dim)))
self.word_emb.weight[self.EOS_IDX].data.copy_(
torch.tensor(np.random.randn(self.emb_dim)))
if self.freeze_embeddings:
self.emb_grad_mask = np.zeros(self.word_emb.weight.shape,
dtype=np.float32)
self.emb_grad_mask[self.UNK_IDX, :] = 1
self.emb_grad_mask[self.PAD_IDX, :] = 1
self.emb_grad_mask[self.START_IDX, :] = 1
self.emb_grad_mask[self.EOS_IDX, :] = 1
self.emb_grad_mask = to_gpu(torch.tensor(self.emb_grad_mask))
else:
self.emb_grad_mask = None
if decode_with_embeddings:
self.lm = nn.GRU(self.emb_dim, self.h_dim, dropout=rnn_dropout)
self.lm_fc = nn.Linear(self.h_dim, self.emb_dim)
else:
self.lm = nn.GRU(self.emb_dim, self.h_dim, dropout=rnn_dropout)
self.lm_fc = nn.Linear(self.h_dim, n_vocab)
self.lm_params = filter(lambda p: p.requires_grad,
self.lm.parameters())
"""
Use GPU if set
"""
if self.gpu:
self.cuda()
def __init__(self, n_vocab, h_dim=128, z_dim=128, p_word_dropout=0.3,
unk_idx=0, pad_idx=1, start_idx=2, eos_idx=3, max_sent_len=15,
use_input_embeddings=True, set_other_to_random=False,
set_unk_to_random=True, decode_with_embeddings=True,
rnn_dropout=0.3, mask_pad=True,
pretrained_embeddings=None, freeze_embeddings=False, gpu=False):
super(RNN_VAE, self).__init__()
self.PAD_IDX = pad_idx
self.UNK_IDX = unk_idx
self.START_IDX = start_idx
self.EOS_IDX = eos_idx
self.MAX_SENT_LEN = max_sent_len
self.freeze_embeddings = freeze_embeddings
self.n_vocab = n_vocab
self.h_dim = h_dim
self.z_dim = z_dim
self.p_word_dropout = p_word_dropout
self.use_input_embeddings = use_input_embeddings # in case we go for sentence embeddings
self.decode_with_embeddings = decode_with_embeddings
self.gpu = gpu
if mask_pad:
self.weights = np.ones(n_vocab)
self.weights[self.PAD_IDX] = 0
self.weights = to_gpu(torch.tensor(self.weights.astype(np.float32)))
else:
self.weights = None
"""
Word embeddings layer
"""
if self.use_input_embeddings:
if pretrained_embeddings is None:
self.emb_dim = h_dim
self.word_emb = nn.Embedding(n_vocab, h_dim, self.PAD_IDX)
else:
self.emb_dim = pretrained_embeddings.size(1)
self.word_emb = nn.Embedding(n_vocab, self.emb_dim, self.PAD_IDX)
# Set pretrained embeddings
self.word_emb.weight.data.copy_(pretrained_embeddings)
else:
raise NotImplementedError('not supported yet')
if set_unk_to_random:
self.word_emb.weight[self.UNK_IDX].data.copy_(torch.tensor(np.random.randn(self.emb_dim)))
if set_other_to_random:
self.word_emb.weight[self.PAD_IDX].data.copy_(torch.tensor(np.random.randn(self.emb_dim)))
self.word_emb.weight[self.START_IDX].data.copy_(torch.tensor(np.random.randn(self.emb_dim)))
self.word_emb.weight[self.EOS_IDX].data.copy_(torch.tensor(np.random.randn(self.emb_dim)))
if self.freeze_embeddings:
self.emb_grad_mask = np.zeros(self.word_emb.weight.shape, dtype=np.float32)
self.emb_grad_mask[self.UNK_IDX,:] = 1
self.emb_grad_mask[self.PAD_IDX,:] = 1
self.emb_grad_mask[self.START_IDX,:] = 1
self.emb_grad_mask[self.EOS_IDX,:] = 1
self.emb_grad_mask = to_gpu(torch.tensor(self.emb_grad_mask))
else:
self.emb_grad_mask = None
"""
Encoder is GRU with FC layers connected to last hidden unit
"""
self.encoder = nn.GRU(self.emb_dim, h_dim)
self.q_mu = nn.Linear(h_dim, z_dim)
self.q_logvar = nn.Linear(h_dim, z_dim)
"""
Decoder is GRU with `z` appended at its inputs
"""
if decode_with_embeddings:
self.decoder = nn.GRU(self.emb_dim+z_dim, z_dim, dropout=rnn_dropout)
self.decoder_fc = nn.Linear(z_dim, self.emb_dim)
else:
self.decoder = nn.GRU(self.emb_dim+z_dim, z_dim, dropout=rnn_dropout)
self.decoder_fc = nn.Linear(z_dim, n_vocab)
"""
Grouping the model's parameters: separating encoder, decoder, and discriminator
"""
self.encoder_params = chain(
self.encoder.parameters(), self.q_mu.parameters(),
self.q_logvar.parameters()
)
self.decoder_params = chain(
self.decoder.parameters(), self.decoder_fc.parameters()
)
self.vae_params = chain(
self.word_emb.parameters(), self.encoder_params, self.decoder_params
)
self.vae_params = filter(lambda p: p.requires_grad, self.vae_params)
"""
Use GPU if set
"""
if self.gpu:
self.cuda()
def sample_z_prior(self, mbsize):
"""
Sample z ~ p(z) = N(0, I)
"""
z = Variable(torch.randn(mbsize, self.z_dim))
return to_gpu(z)
