I’m building a transformer-based architecture using TransformerEncoderLayer and TransformerDecoderLayer provided by PyTorch. The model was running properly without any masks. Since the trained model was printing <pad> in the output, I tried passing a src_key_padding_mask and tgt_key_padding_mask.The mask is created by the method make_padding_key_mask() which returns a ByteTensor of shape (N,S) or (N,T). Now model started getting nan outputs from encoder as well as decoder. Is there anything wrong with way I implemented make_padding_key_mask()
class Transformer_VAE(nn.Module):
def __init__(self, head, vocab_size, embedding_size, latent_dim, device = 'cpu', pad_idx = 0, start_idx = 1, end_idx = 2, unk_idx = 3):
super(Transformer_VAE, self).__init__()
self.head = head
self.embedding_size = embedding_size
self.vocab_size = vocab_size
self.latent_dim = latent_dim
self.device = device
self.embed = WordEmbedding(self.vocab_size, self.embedding_size)
self.postional_encoding = PostionalEncoding(embedding_size, device)
self.encoder = nn.TransformerEncoderLayer(self.embedding_size, self.head)
self.decoder = nn.TransformerDecoderLayer(self.embedding_size, self.head)
self.hidden_to_mean = nn.Linear(self.embedding_size, latent_dim)
self.hidden_to_logvar = nn.Linear(self.embedding_size, latent_dim)
self.latent_to_embed = nn.Linear(self.latent_dim, self.embedding_size)
self.out_linear = nn.Linear(self.embedding_size, vocab_size)
self.pad_idx = pad_idx
self.start_idx = start_idx
self.end_idx = end_idx
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def make_padding_key_mask(self, x):
batch_size, max_len = x.size()[:2]
pad_mask = (x != self.pad_idx).expand(batch_size, max_len)
pad_mask.type(torch.ByteTensor)
return pad_mask
def make_memory_mask(self, tgt_sz, src_sz):
mask = (torch.triu(torch.ones(src_sz, tgt_sz)) == 1).transpose(0, 1)
mem_mask = mask.float().masked_fill(mask == 0, 10e-9).masked_fill(mask == 1, float(0.0))
return mem_mask
def make_tgt_mask(self, tgt_len):
tgt_mask = torch.tril(torch.ones((tgt_len, tgt_len)))
if self.device=='gpu':
tgt_mask = tgt_mask.cuda()
#tgt_mask = tgt_mask.view(batch_size, 1, tgt_len, tgt_len)
return tgt_mask
def get_tgt_embedding(self, tgt, pred):
pred = torch.argmax(pred, dim = -1)
tgt_indices = torch.cat([tgt, pred], dim = 0)
tgt_embed = self.embed(tgt_indices)
return tgt_indices, tgt_embed
def decode(self, tgt, z, tgt_mask, tgt_pad_mask):
z = self.latent_to_embed(z)
out = self.decoder(tgt, z, tgt_mask, tgt_key_padding_mask = tgt_pad_mask)
out = self.out_linear(out)
return out
def forward(self, x):
batch_size, maxlen = x.size()[:2]
tgt_mask = self.make_tgt_mask(maxlen)
mem_mask = self.make_memory_mask(maxlen, maxlen)
src_padding_mask = self.make_padding_key_mask(x)
tgt_padding_mask = self.make_padding_key_mask(x)
if self.device == 'gpu':
tgt_mask = tgt_mask.cuda()
src = self.embed(x).view(maxlen, batch_size, self.embedding_size)
src = self.postional_encoding(src)
x = self.encoder(src, src_key_padding_mask = src_padding_mask)
print(torch.sum(torch.isnan(x)))
mean, logvar = self.hidden_to_mean(x), self.hidden_to_logvar(x)
z = self.reparameterize(mean, logvar)
out = self.decode(src, z, tgt_mask, tgt_padding_mask)
return mean, logvar, out