# How to implement a stateful RNN?

If the model is stateless, the cell states are reset at each example. With the stateful model, all the states are propagated to the next example. Stateful RNNs can learn the long sequences.

For example, we would like to remember some time series.

SeedRandom[0];
data = Accumulate@RandomInteger[{-1, 1}, 1000];
X = data[[;; -2]] // Partition[#, 1] &;
Y = data[[2 ;;]];

ListLinePlot[data]


Imagine a network:

net = NetGraph[
{
LongShortTermMemoryLayer[2],
LinearLayer[1]
},
{
NetPort["InitialState"] -> NetPort[1, "State"],
NetPort["InitialCellState"] -> NetPort[1, "CellState"],
NetPort[1, "State"] -> NetPort["FinalState"],
NetPort[1, "CellState"] -> NetPort["FinalCellState"],
1 -> 2
},
"Input" -> {1, 1},
"Output" -> "Scalar"
] // NetInitialize


During training of the network I want to transfer final states to the initial states for the next example.

net@<|"Input" -> {{1}}, "InitialState" -> {0, 0}, "InitialCellState" -> {0, 0}|>


<|"FinalState" -> {0.0266463, -0.182093}, "FinalCellState" -> {0.0507943, -0.664259}, "Output" -> -0.13134|>

In comments Sebastian recommended me to use NetFoldOperator. But I don't know how to do it.

Also NetTrain shuffles the data during training. What to do?

Example code of stateful RNN in Keras:

model = Sequential()
batch_input_shape=(batch_size, tsteps, 1),
return_sequences=True,
stateful=True))
batch_input_shape=(batch_size, tsteps, 1),
return_sequences=False,
stateful=True))
model.compile(loss='mse', optimizer='rmsprop')

for i in range(epochs):
model.fit(X, Y, batch_size=batch_size, nb_epoch=1, shuffle=False)
model.reset_states()


More examples of stateful RNNs in Keras can be found here: http://machinelearningmastery.com/understanding-stateful-lstm-recurrent-neural-networks-python-keras/

• it would be instructive if you could perhaps post the solution to your problem. I think the accepted answer is informative but does not finish your code. – rfrasier Apr 7 '17 at 0:53
• @rfrasier Thank you for your interest to my question. I have no solution right now. I accepted the answer because I don't expect the more better answer than from one of the developers of the neural networks functionality. When I find a solution, I will post it here. – Alexey Golyshev Apr 7 '17 at 4:38
• thanks for your response. That sounds great. It seems you're interested in neural networks applied to time series. I am as well and so more than happy to share info or help where I can. I'll post anything I work on along those lines. Cheers! – rfrasier Apr 7 '17 at 6:30

The keras link doesn't really say what you're claiming it does. Rather, because Keras doesn't natively support variable-length sequences, you are forced to choose a window size, and then Keras allows you to do your own truncated backpropogation through time by hand to approximate non-windowed training.

Mathematica supports automatic bucketing, and hence you can just feed it examples of different lengths and they will be processed properly, without resorting to truncated backpropogation through time. So the entire reason for wanting to do this stuff goes away, and you simply don't need to worry about any of that. Yay! Problem solved!

However, for extremely long sequences you can find yourself running out of RAM. We plan to address this a different way, which is to use what is sometimes called 'mirroring' to automatically trade off memory for speed until the network fits onto your GPU or whatever. That's for a future version, but in any case you don't need to care about this unless you find yourself actually running out of RAM.

Now, I only realized this AFTER I had described how you would go outside of documented functionality to do this kind of thing. Now that still may be useful or interesting to other people, so even though I think the correct answer to the question is: "you don't need to do that with Mathematica", I'll leave in my alternative answer, too.=, which is below. It shows how to create a stateful trainer object that can be used to write your own training loop.

The language itself is immutable, and so the bias of the deep learning framework is towards being state-free. So there's no way to do what you want that is currently that is documented.

However, you can write your own training loop fairly easily. Here's an example to get you started:

Get["NeuralNetworks"];
net = LinearLayer[];
data = {1 -> 1.9, 2 -> 4.1, 3 -> 6.0, 4 -> 8.1};
{trainer, {gen, nmax}} = NetTrain[
net, data, Automatic,
{"$Trainer", "$TrainingGenerator"},
Method -> {"ADAM", "LearningRateSchedule" -> None},
TrainingProgressReporting -> None,
TrainingProgressFunction -> Function["StopTraining"]
];


This will immediately return two undocumented objects: an MXTrainer and a generator function. The generator function gen can be called with values between 1 and nmax to produce data that can be fed to the trainer. You can also synthesize your own data of course. Here's the actual behavior for this example:

In[79]:= batch = gen[1]
Out[79]= <|"Input" -> {4., 1., 3., 2.}, "Output" -> {8.1, 1.9, 6., 4.1}|>


You can then feed that data to the trainer using:

TrainerUpdate[trainer, batch]


This will fire off a batch update (it happens asynchronously on another thread, the function will return immediately).

You can then get the current loss out like this (this will block until the asynch training update is done):

In[17]:= TrainerCurrentLoss[trainer]
Out[17]= <|"OutputLoss" -> 18.1483|>


Obviously, you just have to call these things in a loop. If you're feeling clever you can load up the next batch of data before you call TrainerCurrentLoss, that will keep the GPU as busy as possible (if you're doing GPU training, of course).

As you can see the loss goes down:

In[19]:= Table[
TrainerUpdate[trainer, batch]; TrainerCurrentLoss[trainer] // First,
30]

Out[19]= {16.3051, 16.2481, 16.1912, 16.1344, 16.0778, 16.0214, \
15.9651, 15.909, 15.853, 15.7971, 15.7414, 15.6858, 15.6303, 15.575, \
15.52, 15.4651, 15.4103, 15.3557, 15.3011, 15.2467, 15.1926, 15.1385, \
15.0846, 15.0308, 14.9771, 14.9237, 14.8704, 14.8171, 14.7642, \
14.7113}


Type ?MXTrainer to see more info about MXTrainer, like how to get the final net back.

Obscure point: Note that I manually disabled the learning rate schedule (ADAM doesn't actually use one by default, but SGD does). That's owing to a small bug. You CAN use a schedule as well, but then you have to set NeuralNetworksPrivateNetTrainmaxBatches to an appropriate value manually.

Now, to actually maintain some state somewhere, such that you can get it out again, you will have to construct your own training network (and of course give a string as the third arg to NetTrain so that it uses your loss port), as well as have an output that actually contains the state that you can extract to use as an input for the next batch. To get access to that output, you can ask the trainer for handles to the underlying arrays it allocated for the net's outputs.

All of this makes things slightly trickier. For one thing, to ensure the non-loss output doesn't get optimized away, you have to include it as a loss port, but with a weight of 0. Here's a toy example:

net = NetGraph[
{{NetPort["Input"], NetPort["StateIn"]} -> 1 -> 2 -> 3 -> NetPort["Loss"], 2 -> NetPort["StateOut"]},
"Input" -> "Real", "Target" -> "Real"
]
data=<|"Input" -> {1,2,3,4}, "StateIn" -> {1,2,3,4}, "Target" -> {5,6,8,9}|>;
{trainer, {gen,nmax}} = NetTrain[net, data, {"Loss" -> Scaled[1], "StateOut" -> Scaled[0]}, {"$Trainer", "$TrainingGenerator"},
TrainingProgressReporting -> None, TrainingProgressFunction -> Function["StopTraining"]];


After doing the update, you can get the state output from the network (to use as the state input for the next batch) like this:

stateOut = Lookup[ReleaseHold @ trainer["OutputArrays"], "StateOut"]


that will be an NDArray[...] object, whose contents you can retrieve with Normal.

• It seems that in the NeuralNetworks context in version 11.1, there are a few functions start with MX (MXPlan, MXTrainer), which are not present in the version 11.0. Could you describe a little bit the design principles behind these functions? – xslittlegrass Apr 5 '17 at 19:24
• @xslittlegrass A lot of them are documented if you look at them with ?. But very briefly MXPlans are device- and batchsize- independent unrollings of a Mathematica net. Once a device and batchsize has been chosen a plan can be instantiated into an executor, which owns the actual memory allocations (NDArrays). An MXTrainer contains an executor (which could be a BucketedExecutor, which creates and caches executors for different sequence length buckets), as well as optimizers, and a cached top-level net (for #Net in the callbacks). – Taliesin Beynon Apr 5 '17 at 20:07
• @TaliesinBeynon Thank you very much for your "alternative answer"! – Alexey Golyshev Apr 6 '17 at 5:00