wassname/NALU-pytorch
1
0
Fork 0
mirror of https://github.com/wassname/NALU-pytorch.git synced 2026-08-20 12:00:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

restructuring and renaming

Browse Source
This commit is contained in:
Kevin
2018-08-04 14:30:47 -07:00
parent 940041d06f
commit 3632c231f0
8 changed files with 155 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
failures.py
+1 -1
View File
@@ -57,7 +57,7 @@ def main():
print("Working with {}...".format(non_lin))
mses = []
for i in range(100):
net = MLP(non_lin)
net = MLP(4, 1, 8, 1, non_lin)
optim = torch.optim.Adam(net.parameters(), lr=LEARNING_RATE)
train(net, optim, train_data, NUM_ITERS)
mses.append(test(net, test_data))
function_learning.py
+14 -13
View File
@@ -7,7 +7,7 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from models import MultiLayerNet, MultiLayerNAC, MultiLayerNALU
from models import MLP, NAC, NALU
NORMALIZE = True
NUM_LAYERS = 2
@@ -68,27 +68,27 @@ def main():
save_dir = './results/'
models = [
MultiLayerNet(
'relu6',
MLP(
num_layers=NUM_LAYERS,
in_dim=2,
hidden_dim=HIDDEN_DIM,
out_dim=1
out_dim=1,
activation='relu6',
),
MultiLayerNet(
'none',
MLP(
num_layers=NUM_LAYERS,
in_dim=2,
hidden_dim=HIDDEN_DIM,
out_dim=1
out_dim=1,
activation='none',
),
MultiLayerNAC(
NAC(
num_layers=NUM_LAYERS,
in_dim=2,
hidden_dim=HIDDEN_DIM,
out_dim=1
out_dim=1,
),
MultiLayerNALU(
NALU(
num_layers=NUM_LAYERS,
in_dim=2,
hidden_dim=HIDDEN_DIM,
@@ -110,9 +110,10 @@ def main():
# random model
random_mse = []
for i in range(100):
net = MultiLayerNet(
'relu6', num_layers=NUM_LAYERS,
in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1
net = MLP(
num_layers=NUM_LAYERS, in_dim=2,
hidden_dim=HIDDEN_DIM, out_dim=1,
activation='relu6',
)
mse = test(net, X_test, y_test)
random_mse.append(mse.mean().item())
models/__init__.py
+2 -3
View File
@@ -1,4 +1,3 @@
from .mlp import MLP
from .nac import NAC
from .nalu import NALU
from .models import MultiLayerNet, MultiLayerNAC, MultiLayerNALU
from .nac import NeuralAccumulatorCell, NAC
from .nalu import NeuralArithmeticLogicUnitCell, NALU
models/mlp.py
+46 -11
View File
@@ -5,15 +5,43 @@ from .utils import str2act
class MLP(nn.Module):
def __init__(self, activation, input_dim=1, encoding_dim=8):
"""A Multi-Layer Perceptron (MLP).
Also known as a Fully-Connected Network (FCN). This
implementation assumes that all hidden layers have
the same hidden size and the same activation function.
Attributes:
num_layers: the number of layers in the network.
in_dim: the size of the input sample.
hidden_dim: the size of the hidden layers.
out_dim: the size of the output.
activation: the activation function.
"""
def __init__(self, num_layers, in_dim, hidden_dim, out_dim, activation='relu'):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
self.activation = str2act(activation)
self.i2h = nn.Linear(input_dim, encoding_dim)
self.h2h1 = nn.Linear(encoding_dim, encoding_dim)
self.h2h2 = nn.Linear(encoding_dim, encoding_dim)
self.h2h3 = nn.Linear(encoding_dim, encoding_dim)
self.h2o = nn.Linear(encoding_dim, input_dim)
nonlin = True
if self.activation is None:
nonlin = False
layers = []
for i in range(num_layers - 1):
layers.extend(
self._layer(
hidden_dim if i > 0 else in_dim,
hidden_dim,
nonlin,
)
)
layers.extend(self._layer(hidden_dim, out_dim, False))
self.model = nn.Sequential(*layers)
# init
for m in self.modules():
@@ -23,10 +51,17 @@ class MLP(nn.Module):
bound = 1 / math.sqrt(fan_in)
nn.init.uniform_(m.bias, -bound, bound)
def _layer(self, in_dim, out_dim, activation=True):
if activation:
return [
nn.Linear(in_dim, out_dim),
self.activation,
]
else:
return [
nn.Linear(in_dim, out_dim),
]
def forward(self, x):
out = self.activation(self.i2h(x))
out = self.activation(self.h2h1(out))
out = self.activation(self.h2h2(out))
out = self.activation(self.h2h3(out))
out = self.h2o(out)
out = self.model(x)
return out
models/models.py
-91
View File
@@ -1,91 +0,0 @@
import math
import torch
import torch.nn as nn
from .nac import NAC
from .nalu import NALU
from .utils import str2act
class MultiLayerNet(nn.Module):
def __init__(self, activation, num_layers, in_dim, hidden_dim, out_dim):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
self.activation = str2act(activation)
layers = []
if self.activation is not None:
layers.extend([
nn.Linear(in_dim, hidden_dim),
self.activation,
])
else:
layers.append(nn.Linear(in_dim, hidden_dim))
for i in range(num_layers - 2):
if self.activation is not None:
layers.extend([
nn.Linear(hidden_dim, hidden_dim),
self.activation,
])
else:
layers.append(nn.Linear(hidden_dim, hidden_dim))
layers.append(nn.Linear(hidden_dim, out_dim))
self.model = nn.Sequential(*layers)
# init
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight)
bound = 1 / math.sqrt(fan_in)
nn.init.uniform_(m.bias, -bound, bound)
def forward(self, x):
out = self.model(x)
return out
class MultiLayerNAC(nn.Module):
def __init__(self, num_layers, in_dim, hidden_dim, out_dim):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
layers = []
layers.append(NAC(in_dim, hidden_dim))
for i in range(num_layers - 2):
layers.append(NAC(hidden_dim, hidden_dim))
layers.append(NAC(hidden_dim, out_dim))
self.model = nn.Sequential(*layers)
def forward(self, x):
out = self.model(x)
return out
class MultiLayerNALU(nn.Module):
def __init__(self, num_layers, in_dim, hidden_dim, out_dim):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
layers = []
layers.append(NALU(in_dim, hidden_dim))
for i in range(num_layers - 2):
layers.append(NALU(hidden_dim, hidden_dim))
layers.append(NALU(hidden_dim, out_dim))
self.model = nn.Sequential(*layers)
def forward(self, x):
out = self.model(x)
return out
models/nac.py
+42 -15
View File
@@ -7,27 +7,23 @@ import torch.nn.functional as F
from torch.nn.parameter import Parameter
class NAC(nn.Module):
"""A Neural Accumulator [1].
NAC supports the ability to accumulate quantities
additively which is a desirable inductive bias for
linear extrapolation.
class NeuralAccumulatorCell(nn.Module):
"""A Neural Accumulator (NAC) cell [1].
Attributes:
in_features: size of the input sample.
out_features: size of the output sample.
in_dim: size of the input sample.
out_dim: size of the output sample.
Sources:
[1]: https://arxiv.org/abs/1808.00508
"""
def __init__(self, in_features, out_features):
def __init__(self, in_dim, out_dim):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.in_dim = in_dim
self.out_dim = out_dim
self.W_hat = Parameter(torch.Tensor(out_features, in_features))
self.M_hat = Parameter(torch.Tensor(out_features, in_features))
self.W_hat = Parameter(torch.Tensor(out_dim, in_dim))
self.M_hat = Parameter(torch.Tensor(out_dim, in_dim))
self.W = Parameter(F.tanh(self.W_hat) * F.sigmoid(self.M_hat))
self.register_parameter('bias', None)
@@ -38,6 +34,37 @@ class NAC(nn.Module):
return F.linear(input, self.W, self.bias)
def extra_repr(self):
return 'in_features={}, out_features={}'.format(
self.in_features, self.out_features
return 'in_dim={}, out_dim={}'.format(
self.in_dim, self.out_dim
)
class NAC(nn.Module):
"""A stack of NAC layers.
Attributes:
num_layers: the number of NAC layers.
in_dim: the size of the input sample.
hidden_dim: the size of the hidden layers.
out_dim: the size of the output.
"""
def __init__(self, num_layers, in_dim, hidden_dim, out_dim):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
layers = []
for i in range(num_layers):
layers.append(
NeuralAccumulatorCell(
hidden_dim if i > 0 else in_dim,
hidden_dim if i < num_layers - 1 else out_dim,
)
)
self.model = nn.Sequential(*layers)
def forward(self, x):
out = self.model(x)
return out
models/nalu.py
+44 -16
View File
@@ -4,33 +4,30 @@ import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from .nac import NAC
from .nac import NeuralAccumulatorCell
from torch.nn.parameter import Parameter
class NALU(nn.Module):
"""A Neural Arithmetic Logic Unit [1].
NALU uses 2 NACs with tied weights to support
multiplicative extrapolation.
class NeuralArithmeticLogicUnitCell(nn.Module):
"""A Neural Arithmetic Logic Unit (NALU) cell [1].
Attributes:
in_features: size of the input sample.
out_features: size of the output sample.
in_dim: size of the input sample.
out_dim: size of the output sample.
Sources:
[1]: https://arxiv.org/abs/1808.00508
"""
def __init__(self, in_features, out_features):
def __init__(self, in_dim, out_dim):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.in_dim = in_dim
self.out_dim = out_dim
self.eps = 1e-10
self.G = Parameter(torch.Tensor(out_features, in_features))
self.W = Parameter(torch.Tensor(out_features, in_features))
self.G = Parameter(torch.Tensor(out_dim, in_dim))
self.W = Parameter(torch.Tensor(out_dim, in_dim))
self.register_parameter('bias', None)
self.nac = NAC(in_features, out_features)
self.nac = NeuralAccumulatorCell(in_dim, out_dim)
init.kaiming_uniform_(self.G, a=math.sqrt(5))
init.kaiming_uniform_(self.W, a=math.sqrt(5))
@@ -46,6 +43,37 @@ class NALU(nn.Module):
return y
def extra_repr(self):
return 'in_features={}, out_features={}'.format(
self.in_features, self.out_features
return 'in_dim={}, out_dim={}'.format(
self.in_dim, self.out_dim
)
class NALU(nn.Module):
"""A stack of NAC layers.
Attributes:
num_layers: the number of NAC layers.
in_dim: the size of the input sample.
hidden_dim: the size of the hidden layers.
out_dim: the size of the output.
"""
def __init__(self, num_layers, in_dim, hidden_dim, out_dim):
super().__init__()
self.num_layers = num_layers
self.in_dim = in_dim
self.hidden_dim = hidden_dim
self.out_dim = out_dim
layers = []
for i in range(num_layers):
layers.append(
NeuralArithmeticLogicUnitCell(
hidden_dim if i > 0 else in_dim,
hidden_dim if i < num_layers - 1 else out_dim,
)
)
self.model = nn.Sequential(*layers)
def forward(self, x):
out = self.model(x)
return out
results/interpolation.txt
+6 -6
View File
@@ -1,7 +1,7 @@
Relu6 None NAC NALU
4.472 0.132 0.154 0.157
85.727 2.224 2.403 34.610
89.257 4.573 5.382 1.236
97.070 60.594 5.730 3.042
89.987 2.977 4.718 1.117
5.939 40.243 7.263 1.119
4.408 0.148 100.085 0.614
93.842 1.937 93.068 58.833
51.517 0.550 100.011 1.448
113.406 3.299 49.287 0.644
48.252 0.735 100.001 1.401
16.141 2.248 98.166 8.952