Hi everyone,
For a personal project I’m trying to recreate the NBeats architecture into Keras, and I don’t think I’m doing it correctly but am not sure why.
The page I’m working off of as a ground truth can be found here: https://github.com/ElementAI/N-BEATS/blob/master/models/nbeats.py
Here’s the starter PyTorch code that I’m trying to convert:
class NBeatsBlock(t.nn.Module):
def __init__(self,
input_size,
theta_size: int,
basis_function: t.nn.Module,
layers: int,
layer_size: int):
super().__init__()
self.layers = t.nn.ModuleList([t.nn.Linear(in_features=input_size, out_features=layer_size)] +
[t.nn.Linear(in_features=layer_size, out_features=layer_size)
for _ in range(layers - 1)])
self.basis_parameters = t.nn.Linear(in_features=layer_size, out_features=theta_size)
self.basis_function = basis_function
def forward(self, x: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:
block_input = x
for layer in self.layers:
block_input = t.relu(layer(block_input))
basis_parameters = self.basis_parameters(block_input)
return self.basis_function(basis_parameters)
class NBeats(t.nn.Module):
def __init__(self, blocks: t.nn.ModuleList):
super().__init__()
self.blocks = blocks
def forward(self, x: t.Tensor, input_mask: t.Tensor) -> t.Tensor:
residuals = x.flip(dims=(1,))
input_mask = input_mask.flip(dims=(1,))
forecast = x[:, -1:]
for i, block in enumerate(self.blocks):
backcast, block_forecast = block(residuals)
residuals = (residuals - backcast) * input_mask
forecast = forecast + block_forecast
return forecast
class GenericBasis(t.nn.Module):
def __init__(self, backcast_size: int, forecast_size: int):
super().__init__()
self.backcast_size = backcast_size
self.forecast_size = forecast_size
def forward(self, theta: t.Tensor):
return theta[:, :self.backcast_size], theta[:, -self.forecast_size:]
Here’s the Keras code I have to translate:
class NBeatsBlock(keras.layers.Layer):
def __init__(self,
theta_size: int,
basis_function: keras.layers.Layer,
layer_size: int = 4):
super(NBeatsBlock, self).__init__()
self.layers_ = [keras.layers.Dense(layer_size, activation = 'relu')
for i in range(layer_size)]
self.basis_parameters = keras.layers.Dense(theta_size)
self.basis_function = basis_function
def call(self, inputs):
x = self.layers_[0](inputs)
for layer in self.layers_[1:]:
x = layer(x)
x = self.basis_parameters(x)
return self.basis_function(x)
class NBeats(keras.layers.Layer):
def __init__(self,
blocksize: int,
theta_size: int,
basis_function: keras.layers.Layer):
super(NBeats, self).__init__()
self.blocks = [NBeatsBlock(theta_size = theta_size, basis_function = basis_function) for i in range(blocksize)]
def call(self, inputs):
residuals = K.reverse(inputs, axes = 0)
forecast = inputs[:, -1:]
for block in self.blocks:
backcast, block_forecast = block(residuals)
residuals = residuals - backcast
forecast = forecast + block_forecast
return forecast
class GenericBasis(keras.layers.Layer):
def __init__(self, backcast_size: int, forecast_size: int):
super().__init__()
self.backcast_size = backcast_size
self.forecast_size = forecast_size
def call(self, inputs):
return inputs[:, :self.backcast_size], inputs[:, -self.forecast_size:]
If I try and make a model from the Keras code it works, but I don’t think it’s constructed correctly.
Here’s a simple model:
inputs = Input(shape = (1, ))
nbeats = NBeats(blocksize = 4, theta_size = 7, basis_function = GenericBasis(7, 7))(inputs)
out = keras.layers.Dense(7)(nbeats)
model = Model(inputs, out)
My concern is that the internal NBeatsBlock layers are not actually being used in the model I just created.
My model summary reads like this:
,
And as you can see there’s nothing that indicates the internal Dense layers are there.
And if I plot the model I get the following diagram:
So I don’t think I’m doing things correctly but I’m also not sure where I’m going wrong with how I’m constructing it. I’m guessing there are small differences in how PyTorch & Keras work that I’m not picking up on.