-
Notifications
You must be signed in to change notification settings - Fork 1
/
pytorch_helper.py
52 lines (37 loc) · 1.57 KB
/
pytorch_helper.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
import torch.nn as nn
def get_init_function(init_value):
def init_function(m):
if init_value > 0.:
if hasattr(m, 'weight'):
m.weight.data.uniform_(-init_value, init_value)
if hasattr(m, 'bias'):
m.bias.data.fill_(0.)
return init_function
class FF(nn.Module):
def __init__(self, dim_input, dim_hidden, dim_output, num_layers,
activation='relu', dropout_rate=0, layer_norm=False,
residual_connection=False):
super().__init__()
assert num_layers >= 0 # 0 = Linear
if num_layers > 0:
assert dim_hidden > 0
if residual_connection:
assert dim_hidden == dim_input
self.residual_connection = residual_connection
self.stack = nn.ModuleList()
for l in range(num_layers):
layer = []
if layer_norm:
layer.append(nn.LayerNorm(dim_input if l == 0 else dim_hidden))
layer.append(nn.Linear(dim_input if l == 0 else dim_hidden,
dim_hidden))
layer.append({'tanh': nn.Tanh(), 'relu': nn.ReLU()}[activation])
if dropout_rate > 0:
layer.append(nn.Dropout(dropout_rate))
self.stack.append(nn.Sequential(*layer))
self.out = nn.Linear(dim_input if num_layers < 1 else dim_hidden,
dim_output)
def forward(self, x):
for layer in self.stack:
x = x + layer(x) if self.residual_connection else layer(x)
return self.out(x)