End-to-end memory modules for machine learning developers, written in Ivy.

• Overview
• Run Through
• Interactive Demos
• Get Involved

What is Ivy Memory?

Ivy memory provides differentiable memory modules, including learnt modules such as Neural Turing Machines (NTM), but also parameter-free modules such as End-to-End Egospheric Spatial Memory (ESM). Check out the docs for more info!

The library is built on top of the Ivy machine learning framework. This means all memory modules simultaneously support: Jax, Tensorflow, PyTorch, MXNet, and Numpy.

Ivy Libraries

There are a host of derived libraries written in Ivy, in the areas of mechanics, 3D vision, robotics, gym environments, neural memory, pre-trained models + implementations, and builder tools with trainers, data loaders and more. Click on the icons below to learn more!

Quick Start

Ivy memory can be installed like so: pip install ivy-memory==0.0.1.post0

To quickly see the different aspects of the library, we suggest you check out the demos! We suggest you start by running the script run_through.py, and read the "Run Through" section below which explains this script.

For more interactive demos, we suggest you run either learning_to_copy_with_ntm.py or mapping_a_room_with_esm.py in the interactive demos folder.

We run through some of the different parts of the library via a simple ongoing example script. The full script is available in the demos folder, as file run_through.py.

End-to-End Egospheric Spatial Memory

First, we show how the Ivy End-to-End Egospheric Spatial Memory (ESM) class can be used inside a pure-Ivy model. We first define the model as below.

class IvyModelWithESM(ivy.Module):

    def __init__(self, channels_in, channels_out):
        self._channels_in = channels_in
        self._esm = ivy_mem.ESM(omni_image_dims=(16, 32))
        self._linear = ivy_mem.Linear(channels_in, channels_out)
        ivy.Module.__init__(self, 'cpu')

    def _forward(self, obs):
        mem = self._esm(obs)
        x = ivy.reshape(mem.mean, (-1, self._channels_in))
        return self._linear(x)

Next, we instantiate this model, and verify that the returned tensors are of the expected shape.

# create model
in_channels = 32
out_channels = 8
ivy.set_backend('torch')
model = IvyModelWithESM(in_channels, out_channels)

# input config
batch_size = 1
image_dims = [5, 5]
num_timesteps = 2
num_feature_channels = 3

# create image of pixel co-ordinates
uniform_pixel_coords =\
    ivy_vision.create_uniform_pixel_coords_image(image_dims, [batch_size, num_timesteps])

# define camera measurement
depths = ivy.random_uniform(shape=[batch_size, num_timesteps] + image_dims + [1])
pixel_coords = ivy_vision.depth_to_pixel_coords(depths)
inv_calib_mats = ivy.random_uniform(shape=[batch_size, num_timesteps, 3, 3])
cam_coords = ivy_vision.pixel_to_cam_coords(pixel_coords, inv_calib_mats)[..., 0:3]
features = ivy.random_uniform(shape=[batch_size, num_timesteps] + image_dims + [num_feature_channels])
img_mean = ivy.concat((cam_coords, features), axis=-1)
cam_rel_mat = ivy.eye(4, batch_shape=[batch_size, num_timesteps])[..., 0:3, :]

# place these into an ESM camera measurement container
esm_cam_meas = ESMCamMeasurement(
    img_mean=img_mean,
    cam_rel_mat=cam_rel_mat
)

# define agent pose transformation
agent_rel_mat = ivy.eye(4, batch_shape=[batch_size, num_timesteps])[..., 0:3, :]

# collect together into an ESM observation container
esm_obs = ESMObservation(
    img_meas={'camera_0': esm_cam_meas},
    agent_rel_mat=agent_rel_mat
)

# call model and test output
output = model(esm_obs)
assert output.shape[-1] == out_channels

Finally, we define a dummy loss function, and show how the ESM network can be trained using Ivy functions only.

# define loss function
target = ivy.zeros_like(output)

def loss_fn(v):
    pred = model(esm_obs, v=v)
    return ivy.mean((pred - target) ** 2)

# optimizer
optimizer = ivy.SGD(lr=1e-4)

# train model
print('\ntraining dummy Ivy ESM model...\n')
for i in range(10):
    loss, grads = ivy.execute_with_gradients(loss_fn, model.v)
    model.v = optimizer.step(model.v, grads)
    print('step {}, loss = {}'.format(i, ivy.to_numpy(loss).item()))
print('\ndummy Ivy ESM model trained!\n')
ivy.unset_backend()

Neural Turing Machine

We next show how the Ivy Neural Turing Machine (NTM) class can be used inside a TensorFlow model. First, we define the model as below.

class TfModelWithNTM(tf.keras.Model):

    def __init__(self, channels_in, channels_out):
        tf.keras.Model.__init__(self)
        self._linear = tf.keras.layers.Dense(64)
        memory_size = 4
        memory_vector_dim = 1
        self._ntm = ivy_mem.NTM(
            input_dim=64, output_dim=channels_out, ctrl_output_size=channels_out, ctrl_layers=1,
            memory_size=memory_size, memory_vector_dim=memory_vector_dim, read_head_num=1, write_head_num=1)
        self._assign_variables()

    def _assign_variables(self):
        self._ntm.v.map(
            lambda x, kc: self.add_weight(name=kc, shape=x.shape))
        self.set_weights([ivy.to_numpy(v) for k, v in self._ntm.v.to_iterator()])
        self.trainable_weights_dict = dict()
        for weight in self.trainable_weights:
            self.trainable_weights_dict[weight.name] = weight
        self._ntm.v = self._ntm.v.map(lambda x, kc: self.trainable_weights_dict[kc + ':0'])

    def call(self, x, **kwargs):
        x = self._linear(x)
        return self._ntm(x)

Next, we instantiate this model, and verify that the returned tensors are of the expected shape.

# create model
in_channels = 32
out_channels = 8
ivy.set_backend('tensorflow')
model = TfModelWithNTM(in_channels, out_channels)

# define inputs
batch_shape = [1, 2]
timesteps = 3
input_shape = batch_shape + [timesteps, in_channels]
input_seq = tf.random.uniform(batch_shape + [timesteps, in_channels])

# call model and test output
output_seq = model(input_seq)
assert input_seq.shape[:-1] == output_seq.shape[:-1]
assert input_seq.shape[-1] == in_channels
assert output_seq.shape[-1] == out_channels

Finally, we define a dummy loss function, and show how the NTM can be trained using a native TensorFlow optimizer.

# define loss function
target = tf.zeros_like(output_seq)

def loss_fn():
    pred = model(input_seq)
    return tf.reduce_sum((pred - target) ** 2)

# define optimizer
optimizer = tf.keras.optimizers.Adam(1e-2)

# train model
print('\ntraining dummy TensorFlow NTM model...\n')
for i in range(10):
    with tf.GradientTape() as tape:
        loss = loss_fn()
    grads = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    print('step {}, loss = {}'.format(i, loss))
print('\ndummy TensorFlow NTM model trained!\n')
ivy.unset_framework()

In addition to the run through above, we provide two further demo scripts, which are more visual and interactive.

The scripts for these demos can be found in the interactive demos folder.

Learning to Copy with NTM

The first demo trains a Neural Turing Machine to copy a sequence from one memory bank to another. NTM can overfit to a single copy sequence very quickly, as show in the real-time visualization below.

Mapping a Room with ESM

The second demo creates an egocentric map of a room, from a rotating camera. The raw image observations are shown on the left, and the incrementally constructed omni-directional ESM feature and depth images are shown on the right. While this example only projects color values into the memory, arbitrary neural network features can also be projected, for end-to-end training.

We hope the memory classes in this library are useful to a wide range of machine learning developers. However, there are many more areas of differentiable memory which could be covered by this library.

If there are any particular functions or classes you feel are missing, and your needs are not met by the functions currently on offer, then we are very happy to accept pull requests!

We look forward to working with the community on expanding and improving the Ivy memory library.

@article{lenton2021ivy,
  title={Ivy: Templated deep learning for inter-framework portability},
  author={Lenton, Daniel and Pardo, Fabio and Falck, Fabian and James, Stephen and Clark, Ronald},
  journal={arXiv preprint arXiv:2102.02886},
  year={2021}
}