mojolearn.RNNRegressor#

class mojolearn.RNNRegressor(hidden_size=16, num_layers=1, optimizer='adam', learning_rate=0.001, batch_size=32, max_epochs=10, shuffle=True, random_state=0, optimizer_options=None, nonlinearity='tanh', numeric_mode=None, predict_chunk=4096, lr_schedule=None)[source]#

Bases: _RecurrentRegressor

Elman RNN sequence regressor: recurrent layers under a linear head.

The nonlinearity is ‘tanh’ or ‘relu’. The recurrent layers are PyTorch’s nn.RNN with batch_first=True: the same parameter names, shapes, gate order and initialization (every weight and bias drawn from U(-1/sqrt(hidden_size), 1/sqrt(hidden_size)), the head as nn.Linear), so state_dict() loads into a torch module and back. A linear head reads the top layer’s last hidden state, and training minimizes the mean squared error by backpropagation through time.

X has shape (n_samples, n_timesteps, n_features) and must be float32 (float64 raises TypeError). y is (n_samples,) or (n_samples, n_outputs).

Not supported: bidirectional layers, dropout between layers, proj_size, an initial hidden state, and packed or variable-length sequences.

Parameters:

  • hidden_size (int, default 16) – Size of the hidden state of each recurrent layer.

  • num_layers (int, default 1) – Number of stacked recurrent layers.

  • optimizer ({‘adam’, ‘adamw’, ‘sgd’, ‘rmsprop’, ‘adagrad’, ‘lion’, ‘adamax’, ‘nadam’}, default ‘adam’) – Optimizer, with PyTorch’s update rule and defaults.

  • learning_rate (float, default 1e-3) – Learning rate, used when lr_schedule is None.

  • batch_size (int, default 32) – Minibatch size (capped at n_samples).

  • max_epochs (int, default 10) – Number of passes over the data. max_epochs * n_samples must be below 2**31 - 1.

  • shuffle (bool, default True) – Reshuffle the rows every epoch.

  • random_state (int or None, default 0) – Seed of the weight initialization and the shuffle.

  • optimizer_options (dict or None, default None) – Extra optimizer arguments under their PyTorch names (for example {'betas': (0.9, 0.99), 'weight_decay': 0.01}). Unknown names raise ValueError.

  • nonlinearity ({‘tanh’, ‘relu’}, default ‘tanh’) – Activation of the recurrence, as nn.RNN.

  • numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

  • predict_chunk (int, default 4096) – Rows processed per chunk at prediction time.

  • lr_schedule (object or None, default None) – A schedule such as StepLR, ExponentialLR or OneCycleLR (anything with a one-based lr_at(t)); it sets the learning rate of every optimizer step.

Variables:

  • n_features_in_ (int) – Number of input features.

  • n_outputs_ (int) – Number of head outputs.

  • n_timesteps_ (int) – Sequence length seen in fit.

  • n_iter_ (int) – Number of optimizer steps taken.

  • loss_curve_ (ndarray) – Training loss at every optimizer step.

  • params_ (ndarray) – All parameters as one flat float32 array.

fit(X, y)#

Train the model; returns self.

predict(X)#

Predicted values.

hidden_sequence(X)#

Top layer’s hidden states, (n_samples, n_timesteps, hidden_size), as nn.RNN(..., batch_first=True)(X)[0].

state_dict()#

Parameters as a dict of arrays under PyTorch names.

load_state_dict(state)#

Load parameters from such a dict; returns self.

Examples

>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> X = rng.random((256, 20, 4), dtype=np.float32)
>>> y = X[:, :, 0].sum(axis=1)
>>> model = mojolearn.RNNRegressor(hidden_size=32, max_epochs=5).fit(X, y)
>>> pred = model.predict(X)
__init__(hidden_size=16, num_layers=1, optimizer='adam', learning_rate=0.001, batch_size=32, max_epochs=10, shuffle=True, random_state=0, optimizer_options=None, nonlinearity='tanh', numeric_mode=None, predict_chunk=4096, lr_schedule=None)#

Methods

fit(X, y)

hidden_sequence(X)

The top layer's hidden states h_1..h_T, (n_samples, n_timesteps, hidden_size): nn.LSTM(..., batch_first=True)(X)[0].

load_state_dict(state)

predict(X)

state_dict()

{torch name: array}: nn.LSTM/nn.GRU/nn.RNN names for the recurrent layers, head.weight and head.bias for the nn.Linear.