End-to-end memory modules for machine learning developers, written in Ivy.
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.
classIvyModelWithESM(ivy.Module):
def__init__(self, channels_in, channels_out):
self._channels_in=channels_inself._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))
returnself._linear(x)Next, we instantiate this model, and verify that the returned tensors are of the expected shape.
# create modelin_channels=32out_channels=8ivy.set_backend('torch')
model=IvyModelWithESM(in_channels, out_channels)
# input configbatch_size=1image_dims= [5, 5]
num_timesteps=2num_feature_channels=3# create image of pixel co-ordinatesuniform_pixel_coords=\
ivy_vision.create_uniform_pixel_coords_image(image_dims, [batch_size, num_timesteps])
# define camera measurementdepths=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 containeresm_cam_meas=ESMCamMeasurement(
img_mean=img_mean,
cam_rel_mat=cam_rel_mat
)
# define agent pose transformationagent_rel_mat=ivy.eye(4, batch_shape=[batch_size, num_timesteps])[..., 0:3, :]
# collect together into an ESM observation containeresm_obs=ESMObservation(
img_meas={'camera_0': esm_cam_meas},
agent_rel_mat=agent_rel_mat
)
# call model and test outputoutput=model(esm_obs)
assertoutput.shape[-1] ==out_channelsFinally, we define a dummy loss function, and show how the ESM network can be trained using Ivy functions only.
# define loss functiontarget=ivy.zeros_like(output)
defloss_fn(v):
pred=model(esm_obs, v=v)
returnivy.mean((pred-target) **2)
# optimizeroptimizer=ivy.SGD(lr=1e-4)
# train modelprint('\ntraining dummy Ivy ESM model...\n')
foriinrange(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.
classTfModelWithNTM(tf.keras.Model):
def__init__(self, channels_in, channels_out):
tf.keras.Model.__init__(self)
self._linear=tf.keras.layers.Dense(64)
memory_size=4memory_vector_dim=1self._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(
lambdax, kc: self.add_weight(name=kc, shape=x.shape))
self.set_weights([ivy.to_numpy(v) fork, vinself._ntm.v.to_iterator()])
self.trainable_weights_dict=dict()
forweightinself.trainable_weights:
self.trainable_weights_dict[weight.name] =weightself._ntm.v=self._ntm.v.map(lambdax, kc: self.trainable_weights_dict[kc+':0'])
defcall(self, x, **kwargs):
x=self._linear(x)
returnself._ntm(x)Next, we instantiate this model, and verify that the returned tensors are of the expected shape.
# create modelin_channels=32out_channels=8ivy.set_backend('tensorflow')
model=TfModelWithNTM(in_channels, out_channels)
# define inputsbatch_shape= [1, 2]
timesteps=3input_shape=batch_shape+ [timesteps, in_channels]
input_seq=tf.random.uniform(batch_shape+ [timesteps, in_channels])
# call model and test outputoutput_seq=model(input_seq)
assertinput_seq.shape[:-1] ==output_seq.shape[:-1]
assertinput_seq.shape[-1] ==in_channelsassertoutput_seq.shape[-1] ==out_channelsFinally, we define a dummy loss function, and show how the NTM can be trained using a native TensorFlow optimizer.
# define loss functiontarget=tf.zeros_like(output_seq)
defloss_fn():
pred=model(input_seq)
returntf.reduce_sum((pred-target) **2)
# define optimizeroptimizer=tf.keras.optimizers.Adam(1e-2)
# train modelprint('\ntraining dummy TensorFlow NTM model...\n')
foriinrange(10):
withtf.GradientTape() astape:
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}
}