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
|
import torch.nn as nn
import torch.nn.functional as F
class DummyModel(nn.Module):
def __init__(
self,
num_embeddings: int,
embedding_dim: int,
dense_input_size: int,
dense_output_size: int,
dense_layers_count: int,
sparse: bool,
):
r"""
A dummy model with an EmbeddingBag Layer and Dense Layer.
Args:
num_embeddings (int): size of the dictionary of embeddings
embedding_dim (int): the size of each embedding vector
dense_input_size (int): size of each input sample
dense_output_size (int): size of each output sample
dense_layers_count: (int): number of dense layers in dense Sequential module
sparse (bool): if True, gradient w.r.t. weight matrix will be a sparse tensor
"""
super().__init__()
self.embedding = nn.EmbeddingBag(num_embeddings, embedding_dim, sparse=sparse)
self.dense = nn.Sequential(
*[
nn.Linear(dense_input_size, dense_output_size)
for _ in range(dense_layers_count)
]
)
def forward(self, x):
x = self.embedding(x)
return F.softmax(self.dense(x), dim=1)
|
