From 3632c231f05fb0aab052b4893a5fd89972a03c28 Mon Sep 17 00:00:00 2001 From: Kevin Date: Sat, 4 Aug 2018 14:30:47 -0700 Subject: [PATCH] restructuring and renaming --- failures.py | 2 +- function_learning.py | 27 ++++++------ models/__init__.py | 5 +-- models/mlp.py | 57 +++++++++++++++++++----- models/models.py | 91 --------------------------------------- models/nac.py | 57 +++++++++++++++++------- models/nalu.py | 60 +++++++++++++++++++------- results/interpolation.txt | 12 +++--- 8 files changed, 155 insertions(+), 156 deletions(-) delete mode 100644 models/models.py diff --git a/failures.py b/failures.py index c001e41..da23bc6 100644 --- a/failures.py +++ b/failures.py @@ -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)) diff --git a/function_learning.py b/function_learning.py index d8ddcfc..3e0e2ba 100644 --- a/function_learning.py +++ b/function_learning.py @@ -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()) diff --git a/models/__init__.py b/models/__init__.py index c9fbfa9..f136986 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -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 diff --git a/models/mlp.py b/models/mlp.py index 7078527..f588911 100644 --- a/models/mlp.py +++ b/models/mlp.py @@ -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 diff --git a/models/models.py b/models/models.py deleted file mode 100644 index e43eea9..0000000 --- a/models/models.py +++ /dev/null @@ -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 diff --git a/models/nac.py b/models/nac.py index c15e74a..b936fda 100644 --- a/models/nac.py +++ b/models/nac.py @@ -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 diff --git a/models/nalu.py b/models/nalu.py index fd4e6f3..65015e4 100644 --- a/models/nalu.py +++ b/models/nalu.py @@ -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 diff --git a/results/interpolation.txt b/results/interpolation.txt index c98c3ee..1f6317b 100644 --- a/results/interpolation.txt +++ b/results/interpolation.txt @@ -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