From ba9fd1a10746d85d2502c8a79ac49db63d346b04 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sun, 9 Apr 2023 11:29:31 +0200 Subject: Update --- models/sparse.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) (limited to 'models/sparse.py') diff --git a/models/sparse.py b/models/sparse.py index d706db5..bcb2897 100644 --- a/models/sparse.py +++ b/models/sparse.py @@ -5,22 +5,29 @@ import torch.nn as nn class PseudoSparseEmbedding(nn.Module): - def __init__(self, embedding_dim: int, device=None, dtype=torch.float32): + def __init__(self, embedding_dim: int, dropout_p: float = 0.0, device=None, dtype=torch.float32): super().__init__() self.embedding_dim = embedding_dim self.dtype = dtype self.params = nn.ParameterList() + + if dropout_p > 0.0: + self.dropout = nn.Dropout(p=dropout_p) + else: + self.dropout = lambda x: x + self.register_buffer('mapping', torch.zeros(0, device=device, dtype=torch.long)) def forward(self, input_ids: torch.LongTensor): - ids = self.mapping[input_ids.to(self.mapping.device)] + input_ids = input_ids.to(self.mapping.device) + ids = self.mapping[input_ids] mask = ~(ids == -1) if torch.all(~mask): embs = None else: - embs = torch.stack([self.params[id] for id in ids[mask]]) + embs = self.dropout(torch.stack([self.params[id] for id in ids[mask]])) return embs, mask -- cgit v1.2.3-54-g00ecf