ubeydemavus / ncps

PyTorch and TensorFlow implementation of NCP, LTC, and CfC wired neural models

Neural Circuit Policies (for PyTorch and TensorFlow)

📜 Papers

Neural Circuit Policies Enabling Auditable Autonomy (Open Access).
Closed-form continuous-time neural networks (Open Access)

Neural Circuit Policies (NCPs) are designed sparse recurrent neural networks loosely inspired by the nervous system of the organism C. elegans. The goal of this package is to making working with NCPs in PyTorch and keras as easy as possible.

📖 Docs

import torch
from ncps.torch import CfC

rnn = CfC(20,50) # (input, hidden units)
x = torch.randn(2, 3, 20) # (batch, time, features)
h0 = torch.zeros(2,50) # (batch, units)
output, hn = rnn(x,h0)

Installation

pip install ncps

🔖 Colab Notebooks

We have created a few Google Colab notebooks for an interactive introduction to the package

• Google Colab (Pytorch) Basic usage
• Google Colab (Tensorflow): Basic usage
• Google Colab (Tensorflow) Stacking NCPs with other layers

End-to-end Examples

• Quickstart (torch and tf)
• Atari Behavior Cloning (torch and tf)
• Atari Reinforcement Learning (tf)

Usage: Models and Wirings

The package provides two models, the liquid time-constant (LTC) and the closed-form continuous-time (CfC) models. Both models are available as tf.keras.layers.Layer or torch.nn.Module RNN layers.

from ncps.torch import CfC, LTC

input_size = 20
units = 28 # 28 neurons
rnn = CfC(input_size, units)
rnn = LTC(input_size, units)

The RNNs defined above consider fully-connected layers, i.e., as in LSTM, GRUs, and other RNNs. The distinctiveness of NCPs is their structured wiring diagram. To combine the LTC or CfC model with a

from ncps.torch import CfC, LTC
from ncps.wirings import AutoNCP

wiring = AutoNCP(28, 4) # 28 neurons, 4 outputs
input_size = 20
rnn = CfC(input_size, wiring)
rnn = LTC(input_size, wiring)

Tensorflow

The Tensorflow bindings are available via the ncps.tf module.

from ncps.tf import CfC, LTC
from ncps.wirings import AutoNCP

units = 28
wiring = AutoNCP(28, 4) # 28 neurons, 4 outputs
input_size = 20
rnn1 = LTC(units) # fully-connected LTC
rnn2 = CfC(units) # fully-connected CfC
rnn3 = LTC(wiring) # NCP wired LTC
rnn4 = CfC(wiring) # NCP wired CfC

We can then combine the NCP cell with arbitrary tf.keras.layers, for instance to build a powerful image sequence classifier:

from ncps.wirings import AutoNCP
from ncps.tf import LTC
import tensorflow as tf
height, width, channels = (78, 200, 3)

ncp = LTC(AutoNCP(32, output_size=8), return_sequences=True)

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.InputLayer(input_shape=(None, height, width, channels)),
        tf.keras.layers.TimeDistributed(
            tf.keras.layers.Conv2D(32, (5, 5), activation="relu")
        ),
        tf.keras.layers.TimeDistributed(tf.keras.layers.MaxPool2D()),
        tf.keras.layers.TimeDistributed(
            tf.keras.layers.Conv2D(64, (5, 5), activation="relu")
        ),
        tf.keras.layers.TimeDistributed(tf.keras.layers.MaxPool2D()),
        tf.keras.layers.TimeDistributed(tf.keras.layers.Flatten()),
        tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(32, activation="relu")),
        ncp,
        tf.keras.layers.TimeDistributed(tf.keras.layers.Activation("softmax")),
    ]
)
model.compile(
    optimizer=tf.keras.optimizers.Adam(0.01),
    loss='sparse_categorical_crossentropy',
)

@article{lechner2020neural,
  title={Neural circuit policies enabling auditable autonomy},
  author={Lechner, Mathias and Hasani, Ramin and Amini, Alexander and Henzinger, Thomas A and Rus, Daniela and Grosu, Radu},
  journal={Nature Machine Intelligence},
  volume={2},
  number={10},
  pages={642--652},
  year={2020},
  publisher={Nature Publishing Group}
}