In this post, we will implement the LongShortTermMemoryLayer from scratch. This will work to lay the foundations to implement any recurrent layer variants. We will implement LSTM in 4 parts, each one having an image to guide.
To concretize things, let’s assume that hidden state is a 32 dimensional vector and the input to LSTM (at every time step) is 128 dimensional.
In[]:=
inputDimension=128;hiddenStateDimension=32;
Implementing first part
Implementing first part
These neat images were taken from Christopher Olah’s blog, anyways getting back to its implementation:
In[]:=
lstmCell=NetGraph[{(*1*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*2*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*3*)TotalLayer[],(*4*)ElementwiseLayer[LogisticSigmoid](*f_t*)},{NetPort["old_h"]23,NetPort["Input"]13,34(*f_t*)},"old_h"{hiddenStateDimension}]
Out[]=
NetGraph
Implementing second part
Implementing second part
In[]:=
lstmCell=NetGraph[{(*1*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*2*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*3*)TotalLayer[],(*4*)ElementwiseLayer[LogisticSigmoid](*f_t*),(*5*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*6*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*7*)TotalLayer[],(*8*)ElementwiseLayer[LogisticSigmoid](*i_t*),(*9*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*10*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*11*)TotalLayer[],(*12*)ElementwiseLayer[Tanh](*c_tilde_t*)},{NetPort["Input"]13,NetPort["old_h"]23,34(*f_t*),NetPort["Input"]57,NetPort["old_h"]67,78(*i_t*),NetPort["Input"]911,NetPort["old_h"]1011,1112(*c_tilde_t*)},"old_h"{hiddenStateDimension}]
Out[]=
NetGraph
Implementing third part
Implementing third part
In[]:=
lstmCell=NetGraph[{(*1*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*2*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*3*)TotalLayer[],(*4*)ElementwiseLayer[LogisticSigmoid](*f_t*),(*5*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*6*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*7*)TotalLayer[],(*8*)ElementwiseLayer[LogisticSigmoid](*i_t*),(*9*)LinearLayer[hiddenStateDimension,"Input"inputDimension],(*10*)LinearLayer[hiddenStateDimension,"Input"hiddenStateDimension],(*11*)TotalLayer[],(*12*)ElementwiseLayer[Tanh](*c_tilde_t*),(*13*)ThreadingLayer[Times],(*14*)ThreadingLayer[Times],(*15*)TotalLayer[](*new_c*)},{NetPort["Input"]13,NetPort["old_h"]23,34(*f_t*),NetPort["Input"]57,NetPort["old_h"]67,78(*i_t*),NetPort["Input"]911,NetPort["old_h"]1011,1112(*c_tilde_t*),413,NetPort["old_c"]13,814,1214,1315,1415(*new_c*),15NetPort["new_c"]},"old_h"{hiddenStateDimension},"old_c"{hiddenStateDimension}]
Out[]=
NetGraph
Implementing last part
Implementing last part
Wrapping into a function
Wrapping into a function
Let’s wrap the above code into a function to generate new custom layers as and when needed:
An example
An example
This one is from LongShortTermMemoryLayer documentation but reworked using custom LSTM layer. It takes in a string “1+10” and outputs the result of that computation.
Creating the custom layers:
Creating a training net:
Extracting trained net and adding net encoder:
Getting network results:
Voila, you have a custom RNN layer implemented!
References
References