-
Notifications
You must be signed in to change notification settings - Fork 7
/
blocks.py
executable file
·48 lines (36 loc) · 1.16 KB
/
blocks.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
import torch
import torch.nn as nn
import torch.nn.functional as F
def stick_breaking(logits):
e = F.sigmoid(logits)
z = (1 - e).cumprod(dim=1)
p = torch.cat([e.narrow(1, 0, 1), e[:, 1:] * z[:, :-1]], dim=1)
return p
def softmax(x, mask=None):
max_x, _ = x.max(dim=-1, keepdim=True)
e_x = torch.exp(x - max_x)
if not (mask is None):
e_x = e_x * mask
out = e_x / (e_x.sum(dim=-1, keepdim=True) + 1e-8)
return out
class ResBlock(nn.Module):
def __init__(self, ninp, nout, dropout, nlayers=1):
super(ResBlock, self).__init__()
self.nlayers = nlayers
self.drop = nn.Dropout(dropout)
self.res = nn.ModuleList(
[nn.Sequential(
nn.Linear(ninp, ninp),
nn.BatchNorm1d(ninp),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(ninp, ninp),
nn.BatchNorm1d(ninp),
)
for _ in range(nlayers)]
)
def forward(self, input):
# input = self.drop(input)
for i in range(self.nlayers):
input = F.relu(self.res[i](input) + input)
return input