-
Notifications
You must be signed in to change notification settings - Fork 2
/
nas_embedding.py
48 lines (39 loc) · 1.55 KB
/
nas_embedding.py
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
37
38
39
40
41
42
43
44
45
46
47
48
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Ravi Krishna 07/25/21
# Import statements.
import torch
import torch.nn as nn
import numpy as np
class EmbeddingDLRM(nn.Module):
"""
Almost equivalent to nn.Embedding
(some abilities of nn.Embedding do not exist in EmbeddingDLRM),
except that it uses the weight initialization from
dlrm_s_pytorch.py.
"""
def __init__(self, num_features, embedding_dimension):
# Superclass initialization.
super(EmbeddingDLRM, self).__init__()
# Store for later.
self.num_features = num_features
self.embedding_dimension = embedding_dimension
# Create the embedding.
self.embedding = nn.Embedding(num_features, embedding_dimension)
with torch.no_grad():
# Change embedding initialization to match
# original DLRM code from dlrm_s_pytorch.py.
W = np.random.uniform(low=-np.sqrt(1.0 / float(self.num_features)),
high=np.sqrt(1.0 / float(self.num_features)),
size=(self.num_features,
self.embedding_dimension)).astype(np.float32)
self.embedding.weight.data = torch.tensor(W, requires_grad=True)
def forward(self, input_indices):
"""
Run input_indices through the embedding.
"""
return self.embedding(input_indices)