| """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) | |