File size: 691 Bytes
fd448dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | """Embedding layers used by TaoNet."""
import torch.nn as nn
class FactorizedEmbedding(nn.Module):
"""Low-rank factorized embedding."""
def __init__(self, vocab_size, d_model, d_embed_rank=96):
super().__init__()
self.vocab_size = vocab_size
self.d_model = d_model
self.d_embed_rank = d_embed_rank
self.embed = nn.Embedding(vocab_size, d_embed_rank)
self.proj = nn.Linear(d_embed_rank, d_model, bias=False)
nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
nn.init.normal_(self.proj.weight, mean=0.0, std=0.02)
def forward(self, input_ids):
x = self.embed(input_ids)
return self.proj(x)
|