Spaces:
Runtime error
Runtime error
File size: 4,127 Bytes
c5ca37a |
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 |
from itertools import chain
import math
import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
from .gaussian_encoder import GaussianEncoderBase
from ..utils import log_sum_exp
class GaussianLSTMEncoder(GaussianEncoderBase):
"""Gaussian LSTM Encoder with constant-length input"""
def __init__(self, args, vocab_size, model_init, emb_init):
super(GaussianLSTMEncoder, self).__init__()
self.ni = args.ni
self.nh = args.enc_nh
self.nz = args.nz
self.args = args
self.embed = nn.Embedding(vocab_size, args.ni)
self.lstm = nn.LSTM(input_size=args.ni,
hidden_size=args.enc_nh,
num_layers=1,
batch_first=True,
dropout=0)
self.linear = nn.Linear(args.enc_nh, 2 * args.nz, bias=False)
self.reset_parameters(model_init, emb_init)
def reset_parameters(self, model_init, emb_init):
# for name, param in self.lstm.named_parameters():
# # self.initializer(param)
# if 'bias' in name:
# nn.init.constant_(param, 0.0)
# # model_init(param)
# elif 'weight' in name:
# model_init(param)
# model_init(self.linear.weight)
# emb_init(self.embed.weight)
for param in self.parameters():
model_init(param)
emb_init(self.embed.weight)
def forward(self, input):
"""
Args:
x: (batch_size, seq_len)
Returns: Tensor1, Tensor2
Tensor1: the mean tensor, shape (batch, nz)
Tensor2: the logvar tensor, shape (batch, nz)
"""
# (batch_size, seq_len-1, args.ni)
word_embed = self.embed(input)
_, (last_state, last_cell) = self.lstm(word_embed)
mean, logvar = self.linear(last_state).chunk(2, -1)
# fix variance as a pre-defined value
if self.args.fix_var > 0:
logvar = mean.new_tensor([[[math.log(self.args.fix_var)]]]).expand_as(mean)
return mean.squeeze(0), logvar.squeeze(0)
# def eval_inference_mode(self, x):
# """compute the mode points in the inference distribution
# (in Gaussian case)
# Returns: Tensor
# Tensor: the posterior mode points with shape (*, nz)
# """
# # (batch_size, nz)
# mu, logvar = self.forward(x)
class VarLSTMEncoder(GaussianLSTMEncoder):
"""Gaussian LSTM Encoder with variable-length input"""
def __init__(self, args, vocab_size, model_init, emb_init):
super(VarLSTMEncoder, self).__init__(args, vocab_size, model_init, emb_init)
def forward(self, input):
"""
Args:
input: tuple which contains x and sents_len
x: (batch_size, seq_len)
sents_len: long tensor of sentence lengths
Returns: Tensor1, Tensor2
Tensor1: the mean tensor, shape (batch, nz)
Tensor2: the logvar tensor, shape (batch, nz)
"""
input, sents_len = input
# (batch_size, seq_len, args.ni)
word_embed = self.embed(input)
packed_embed = pack_padded_sequence(word_embed, sents_len.tolist(), batch_first=True)
_, (last_state, last_cell) = self.lstm(packed_embed)
mean, logvar = self.linear(last_state).chunk(2, -1)
return mean.squeeze(0), logvar.squeeze(0)
def encode(self, input, nsamples):
"""perform the encoding and compute the KL term
Args:
input: tuple which contains x and sents_len
Returns: Tensor1, Tensor2
Tensor1: the tensor latent z with shape [batch, nsamples, nz]
Tensor2: the tenor of KL for each x with shape [batch]
"""
# (batch_size, nz)
mu, logvar = self.forward(input)
# (batch, nsamples, nz)
z = self.reparameterize(mu, logvar, nsamples)
KL = 0.5 * (mu.pow(2) + logvar.exp() - logvar - 1).sum(dim=1)
return z, KL
|