Package: ivy-memory
End-to-end memory modules for machine learning developers, written in Ivy.ContentsOverviewRun ThroughInteractive DemosGet InvoledOverviewWhat 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 thedocsfor 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 LibrariesThere 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 StartIvy memory can be installed like so:pip installivy-memoryTo quickly see the different aspects of the library, we suggest you check out thedemos!
We suggest you start by running the scriptrun_through.py,
and read the “Run Through” section below which explains this script.For more interactive demos, we suggest you run eitherlearning_to_copy_with_ntm.pyormapping_a_room_with_esm.pyin theinteractivedemos folder.Run ThroughWe run through some of the different parts of the library via a simple ongoing example script.
The full script is available in thedemosfolder, as filerun_through.py.End-to-End Egospheric Spatial MemoryFirst, 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_framework('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.concatenate((cam_coords,features),-1)cam_rel_mat=ivy.identity(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.identity(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.reduce_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_framework()Neural Turing MachineWe 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_framework('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()Interactive DemosIn 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 theinteractivedemos folder.Learning to Copy with NTMThe 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 ESMThe 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.Get InvoledWe 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.Citation@article{lenton2021ivy,
  title={Ivy: Unified Machine 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}
}