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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
import torch.nn as nn
import torch as th
from model.nn.eprop_gate_l0rd import EpropGateL0rd
from model.nn.eprop_transformer_utils import AlphaAttention, InputEmbeding, OutputEmbeding
class EpropGateL0rdTransformer(nn.Module):
def __init__(
self,
channels,
multiplier,
num_objects,
batch_size,
heads,
depth,
reg_lambda,
dropout=0.0
):
super(EpropGateL0rdTransformer, self).__init__()
num_inputs = channels
num_outputs = channels
num_hidden = channels
num_hidden = channels * multiplier
print(f"Predictor channels: {num_hidden}@({num_hidden // heads}x{heads})")
self.depth = depth
_layers = []
_layers.append(InputEmbeding(num_inputs, num_hidden))
for i in range(depth):
_layers.append(AlphaAttention(num_hidden, num_objects, heads, dropout))
_layers.append(EpropAlphaGateL0rd(num_hidden, batch_size * num_objects, reg_lambda))
_layers.append(OutputEmbeding(num_hidden, num_outputs))
self.layers = nn.Sequential(*_layers)
self.attention = []
self.l0rds = []
for l in self.layers:
if 'AlphaAttention' in type(l).__name__:
self.attention.append(l)
elif 'EpropAlphaGateL0rd' in type(l).__name__:
self.l0rds.append(l)
def get_openings(self):
openings = []
for i in range(self.depth):
openings.append(self.l0rds[i].l0rd.openings_perslot)
openings = th.mean(th.stack(openings, dim=0), dim=0)
return openings
def get_hidden(self):
states = []
for i in range(self.depth):
states.append(self.layers[2 * (i + 1)].l0rd.get_hidden())
return th.cat(states, dim=1)
def set_hidden(self, hidden):
states = th.chunk(hidden, self.depth, dim=1)
for i in range(self.depth):
self.layers[2 * (i + 1)].l0rd.set_hidden(states[i])
def forward(self, input: th.Tensor) -> th.Tensor:
return self.layers(input)
class EpropAlphaGateL0rd(nn.Module):
def __init__(self, num_hidden, batch_size, reg_lambda):
super(EpropAlphaGateL0rd, self).__init__()
self.alpha = nn.Parameter(th.zeros(1)+1e-12)
self.l0rd = EpropGateL0rd(
num_inputs = num_hidden,
num_hidden = num_hidden,
num_outputs = num_hidden,
reg_lambda = reg_lambda,
batch_size = batch_size
)
def forward(self, input):
return input + self.alpha * self.l0rd(input)
|
