-
Notifications
You must be signed in to change notification settings - Fork 2
/
nas_linear.py
68 lines (52 loc) · 1.91 KB
/
nas_linear.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
#!/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 LinearDLRM(nn.Module):
"""
This layer is essentially equivalent
to a regular nn.Linear layer, except
that it uses the weight initialization
from dlrm_s_pytorch.py.
"""
def __init__(self,
in_feat,
out_feat,
bias=False):
# Superclass initialization.
super(LinearDLRM, self).__init__()
# Store for later.
self.in_feat = in_feat
self.out_feat = out_feat
self.bias = bias
# Create the layer.
self.linear_layer = nn.Linear(self.in_feat,
self.out_feat,
bias=self.bias)
# Change initialization to match DLRM code
# from dlrm_s_pytorch.py.
with torch.no_grad():
mean = 0.0
std_dev = np.sqrt(2.0 / (float(self.in_feat) + float(self.out_feat)))
W = np.random.normal(mean,
std_dev,
size=(self.out_feat, self.in_feat)).astype(np.float32)
std_dev = np.sqrt(1.0 / float(self.out_feat))
bt = np.random.normal(mean,
std_dev,
size=self.out_feat).astype(np.float32)
self.linear_layer.weight.data = torch.tensor(W, requires_grad=True)
if self.bias is True:
self.linear_layer.bias.data = torch.tensor(bt, requires_grad=True)
def forward(self, input_vectors):
"""
Run input_vectors through the layer.
"""
return self.linear_layer(input_vectors)