```python
import numpy as np
from collections import OrderedDict
import tensorflow as tf
from tensorflow.contrib.eager.python import tfe

# This silent device placement will hurt performance on GPUs, but we are sticking to CPUs only so it
# should be fine
tf.enable_eager_execution()
tf.set_random_seed(0)
```

    D:\Users\Yue\Anaconda3\lib\site-packages\h5py\__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
      from ._conv import register_converters as _register_converters

# Goal
We are going to implement the JANet model from the paper ["The unreasonable effectiveness of the forget gate"](https://arxiv.org/abs/1804.04849) and try the addition experiment to see if it converges properly.

This is not going to be a high performance model, nor will it probably match the performance of the paper. This is going to be a simple exercise in how easy it is to write down research models in Tensorflow Eager without all the baggage of Tensorflow graphs.

# Components we need
For JANet, we need two components :

1) Chrono Initializer from the paper [Can recurrent neural networks warp time? ](https://openreview.net/forum?id=SJcKhk-Ab)

2) The RNN cell which will do most of the work

Lets create the initializer first since it is simpler

# Chrono Initializer
In short, if we know the length of the sequence in the range $[T_{min}, T_{max}]$, then it would be beneficial to initialize the model with a forgetting time in the same range. Therefore, this would be equivalent to initialize the gate $g$ in the range $[\frac{1}{T_{max}}, \frac{1}{T_{min}}]$.

If the values of both inputs and hidden layers are centered over time, $g(t)$ will typically take values centered around $\sigma(b_g)$. To obtain values in the desired range above, the biases $b_g$ must be in the range $-log(T_{max} - 1)$ and $-log(T_{min} - 1)$.

For the $t$th time step, the forget gate $f_t$ corresponds to $1 - g_t$, where as the input gate $i_t$ corresponds to $g_t$ in an LSTM RNN. Letting the minimum number of timesteps to be 1 and maximum to be $T_{max}$, the paper suggests initializing the forget gate and input gate as follows :

\begin{align}
b_f &\sim log(U([1, T_{max}]) - 1) \\\
b_i &= -b_f
\end{align}

For a complete explanation of ChronoInitializer, refer to the above paper.

```python
if tfe.num_gpus() > 0:
    _device = '/gpu:0'
else:
    _device = '/cpu:0'
```

```python
class ChronoInitializer(tf.keras.initializers.RandomUniform):
    """
    Chrono Initializer from the paper :
    [Can recurrent neural networks warp time?](https://openreview.net/forum?id=SJcKhk-Ab)
    """

    def __init__(self, max_timesteps, seed=None):
        super(ChronoInitializer, self).__init__(1., max_timesteps - 1, seed)
        self.max_timesteps = max_timesteps

    def __call__(self, shape, dtype=None, partition_info=None):
        values = super(ChronoInitializer, self).__call__(shape, dtype=dtype, partition_info=partition_info)
        return tf.log(values)

    def get_config(self):
        config = {
            'max_timesteps': self.max_timesteps
        }
        base_config = super(ChronoInitializer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

```

# RNN Cell

The model described in the paper ["The unreasonable effectiveness of the forget gate"](https://arxiv.org/abs/1804.04849) is called JANet - a model where the forget $(g_f)$ gate alone is the one of two gates, alongside with the memory $(g_c)$ gate. It doesn't use the input $(g_i)$ and output $(g_o)$ gates.

```python

class JANETModel(tf.keras.Model):

    def __init__(self, units, num_outputs, num_timesteps, output_activation='sigmoid', **kwargs):
        super(JANETModel, self).__init__(**kwargs)

        self.units = units
        self.classes = num_outputs
        self.num_timesteps = num_timesteps

        # Initialize the forget gate with ChronoInitializer
        # The memory gate is initialized with zeros
        def bias_initializer(_, *args, **kwargs):
            forget_gate = ChronoInitializer(self.num_timesteps)((self.units,), *args, **kwargs)

            return tf.keras.backend.concatenate([
                forget_gate,
                tf.keras.initializers.Zeros()((self.units,), *args, **kwargs),
            ])

            # Initialize the hidden and recurrent gates
        self.kernel = tf.keras.layers.Dense(2 * units, use_bias=False,
                                            kernel_initializer='glorot_uniform')

        self.recurrent_kernel = tf.keras.layers.Dense(2 * units,
                                                      kernel_initializer='glorot_uniform',
                                                      bias_initializer=bias_initializer)

        # Initialize the final layer (for classification or regression depending on the output activation)
        self.output_dense = tf.keras.layers.Dense(num_outputs, activation=output_activation)

    """
    This is required to override a certain issue with how the weights of a model are loaded
    by a checkpoint.

    While we could use `model.call(inputs)` directly, the more pythonic way of doing this is
    by using `model(inputs)`.

    However, in Eager mode, the ordinary function call does *not* forward the method to the
    models `call` method. This bypasses that issue.
    """
    def __call__(self, *args, **kwargs):
        if not tf.executing_eagerly():
            super(JANETModel, self).__call__(*args, **kwargs)
        return self.call(*args, **kwargs)

    def call(self, inputs, training=None, mask=None):
        # Initialize the hidden memory states
        outputs = []
        states = []

        h_state = tf.zeros((inputs.shape[0], self.units))
        c_state = tf.zeros((inputs.shape[0], self.units))

        # Input is in the shape [None, timesteps, input_dim]
        for t in range(inputs.shape[1]):
            ip = inputs[:, t, :]  # access the T'th timestep

            # Perform the forward pass of the model
            z = self.kernel(ip)
            z += self.recurrent_kernel(h_state)

            # Split the output into the forget and memory outputs
            z0 = z[:, :self.units]
            z1 = z[:, self.units: 2 * self.units]

            # gate updates
            f = tf.keras.activations.sigmoid(z0)
            c = f * c_state + (1. - f) * tf.nn.tanh(z1)

            # state updates
            h = c

            # update our previous state
            h_state = h
            c_state = c

            # preserve the history of our states
            outputs.append(h)
            states.append([h, c])  # here we maintain just 1 state variable, can be more

        # save memory states to global variables so we can access them later
        # self.cell_outputs = tf.stack(outputs, axis=1)
        # self.cell_states = states

        # perform forward pass of the classifier / regressor
        preds = self.output_dense(outputs[-1])

        return preds

```

# Setup
In Eager, if you use only tf.keras layers, then training is a lot simpler. You can simply call the .fit and .predict methods after compiling the models. This is what we will not do for the current model, even though we could, as it is important to explain the limitations of sticking to models with only Keras layers.

We will declare some helper methods, along with the dataset loader. Since the addition task is simple enough, we will not be using the canonical way to load data - tf.data pipeline. For this dataset in which we generate synthetic data, the overhead of the tf.data pipeline is more than the benefit.

```python
# we define some constants
import os
if not os.path.exists('weights'):
    os.makedirs('weights/')

# Parameters taken from https://arxiv.org/abs/1804.04849
TIME_STEPS = 100
NUM_UNITS = 128
LEARNING_RATE = 0.001
STEPS_PER_EPOCH = 100
NUM_EPOCHS = 10
BATCH_SIZE = 50

CHECKPOINTS_DIR = 'checkpoints_addition/'
CHECKPOINT_PATH = CHECKPOINTS_DIR + 'addition_janet.ckpt'

```

# Addition Dataset Generator

```python
# Code reused from https://github.com/batzner/indrnn/blob/master/examples/addition_rnn.py
# with some modifications to run on Tensorflow Eager mode

def batch_generator():
    while True:
        """Generate the adding problem dataset"""
        # Build the first sequence
        add_values = np.random.rand(BATCH_SIZE, TIME_STEPS)

        # Build the second sequence with one 1 in each half and 0s otherwise
        add_indices = np.zeros_like(add_values, dtype='float32')
        half = int(TIME_STEPS / 2)
        for i in range(BATCH_SIZE):
            first_half = np.random.randint(half)
            second_half = np.random.randint(half, TIME_STEPS)
            add_indices[i, [first_half, second_half]] = 1.

        # Zip the values and indices in a third dimension:
        # inputs has the shape (batch_size, time_steps, 2)
        inputs = np.dstack((add_values, add_indices))
        targets = np.sum(np.multiply(add_values, add_indices), axis=1)
        targets = np.expand_dims(targets, -1)

        # center at zero mean
        inputs -= np.mean(inputs, axis=0, keepdims=True)

        inputs = tf.constant(inputs, dtype=tf.float32)
        targets = tf.constant(targets, dtype=tf.float32)

        yield inputs, targets

```

# Loss and Gradients
The following functions simply compute the mean squared error, and compute the gradients of the model.

```python
def loss(y_true, y_pred):
    loss_val = tf.losses.mean_squared_error(y_true, y_pred)
    return loss_val

def grad(model, X, y):
    with tf.GradientTape() as tape:
        preds = model(X)
        loss_val = loss(y, preds)

    vars = model.trainable_variables
    grads = tape.gradient(loss_val, vars)
    grad_vars = zip(grads, vars)

    return grad_vars, loss_val
```

# Training and Checkpointing

```python
import shutil

with tf.device(_device):
    model = JANETModel(NUM_UNITS, num_outputs=1, num_timesteps=TIME_STEPS, output_activation='linear')

    optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE)
    global_step = tf.train.get_or_create_global_step()

    generator = batch_generator()

# try using different optimizers and different optimizer configs
# model.compile(loss='mse', optimizer=optimizer)

best_loss = 100.

loss_history = []
update_counter = 1

for epoch in range(NUM_EPOCHS):
    print("Epoch # : ", epoch + 1)

    for step in range(STEPS_PER_EPOCH):
        # get batch dataset
        with tf.device(_device):
            inputs, targets = next(generator)

            # get gradients and loss at this iteration
            gradients, loss_val = grad(model, inputs, targets)

            # apply gradients
            optimizer.apply_gradients(gradients, tf.train.get_or_create_global_step())

        loss_history.append(loss_val.numpy())
        print(update_counter, ":", loss_history[-1])

        update_counter += 1
    print()

    # remove old checkpoint that we no longer need
    if os.path.exists(CHECKPOINTS_DIR):
        shutil.rmtree(CHECKPOINTS_DIR)

    # save the checkpoint weights
    checkpoint = tfe.Checkpoint(model=model).save(CHECKPOINT_PATH)

    # Optional : Save the weight matrices in Keras format as well
    # model.save_weights('addition_model.h5', overwrite=True)

    print("Saving weights")
    print()

```

    Epoch # :  1
    1 : 1.4451375
    2 : 1.2100395
    3 : 1.0465797
    4 : 1.1592366
    5 : 0.89892256
    6 : 0.95416474
    7 : 0.76734453
    8 : 0.8083082
    9 : 0.5987345
    10 : 0.43943503
    11 : 0.36136585
    12 : 0.1923017
    13 : 0.20704773
    14 : 0.31842035
    15 : 0.32538295
    16 : 0.23726974
    17 : 0.2102163
    18 : 0.17351742
    19 : 0.15244175
    20 : 0.17523283
    21 : 0.27364022
    22 : 0.2666938
    23 : 0.21667854
    24 : 0.20420966
    25 : 0.21706091
    26 : 0.13949679
    27 : 0.22169092
    28 : 0.14021274
    29 : 0.13293959
    30 : 0.13264558
    31 : 0.14265671
    32 : 0.18486334
    33 : 0.17616434
    34 : 0.19732943
    35 : 0.18540621
    36 : 0.20976181
    37 : 0.16762428
    38 : 0.15358415
    39 : 0.09938493
    40 : 0.16012174
    41 : 0.1352579
    42 : 0.13615915
    43 : 0.19032782
    44 : 0.21800873
    45 : 0.120117724
    46 : 0.19352394
    47 : 0.12246391
    48 : 0.16440761
    49 : 0.18718316
    50 : 0.15584385
    51 : 0.18077189
    52 : 0.20388669
    53 : 0.16069853
    54 : 0.12863132
    55 : 0.1536328
    56 : 0.14860286
    57 : 0.14560589
    58 : 0.19622646
    59 : 0.123918116
    60 : 0.15879595
    61 : 0.11481209
    62 : 0.17653896
    63 : 0.13256983
    64 : 0.2005058
    65 : 0.21987212
    66 : 0.136102
    67 : 0.18088768
    68 : 0.22262762
    69 : 0.10627134
    70 : 0.2091097
    71 : 0.21229905
    72 : 0.20627868
    73 : 0.1690417
    74 : 0.17177807
    75 : 0.18062381
    76 : 0.13649574
    77 : 0.17777382
    78 : 0.14919841
    79 : 0.1768317
    80 : 0.16898365
    81 : 0.1236533
    82 : 0.17730775
    83 : 0.1408914
    84 : 0.16091518
    85 : 0.14103621
    86 : 0.1576444
    87 : 0.1519246
    88 : 0.13818125
    89 : 0.1296538
    90 : 0.14973374
    91 : 0.13017823
    92 : 0.14200756
    93 : 0.16256008
    94 : 0.18772072
    95 : 0.17444259
    96 : 0.17605019
    97 : 0.20175363
    98 : 0.13228765
    99 : 0.17516196
    100 : 0.16253173

    Saving weights

    Epoch # :  2
    101 : 0.16850589
    102 : 0.09396511
    103 : 0.17106655
    104 : 0.17185669
    105 : 0.15050283
    106 : 0.22392723
    107 : 0.16291136
    108 : 0.15039755
    109 : 0.18366799
    110 : 0.15504017
    111 : 0.14269602
    112 : 0.14168854
    113 : 0.14030021
    114 : 0.127676
    115 : 0.15085882
    116 : 0.15884759
    117 : 0.13181584
    118 : 0.13127232
    119 : 0.14388686
    120 : 0.1577321
    121 : 0.120080836
    122 : 0.20732296
    123 : 0.13947405
    124 : 0.124289535
    125 : 0.21791926
    126 : 0.25348213
    127 : 0.18623595
    128 : 0.13081437
    129 : 0.14789547
    130 : 0.16478452
    131 : 0.15387581
    132 : 0.18875591
    133 : 0.1939753
    134 : 0.22788465
    135 : 0.17664723
    136 : 0.1677759
    137 : 0.18561229
    138 : 0.19522236
    139 : 0.22043419
    140 : 0.16887332
    141 : 0.19686228
    142 : 0.15505163
    143 : 0.16156386
    144 : 0.20846851
    145 : 0.14506882
    146 : 0.0969856
    147 : 0.123437144
    148 : 0.13982366
    149 : 0.20551193
    150 : 0.15449674
    151 : 0.20522118
    152 : 0.13730673
    153 : 0.18335243
    154 : 0.15370485
    155 : 0.12530635
    156 : 0.16562468
    157 : 0.12206404
    158 : 0.12293398
    159 : 0.18548702
    160 : 0.13402353
    161 : 0.17124395
    162 : 0.19007084
    163 : 0.14525324
    164 : 0.14681573
    165 : 0.16420875
    166 : 0.1213709
    167 : 0.1480321
    168 : 0.123643875
    169 : 0.19550744
    170 : 0.15400517
    171 : 0.18444566
    172 : 0.16708651
    173 : 0.20047146
    174 : 0.15286522
    175 : 0.20791475
    176 : 0.16400278
    177 : 0.16670024
    178 : 0.13447633
    179 : 0.13560998
    180 : 0.16651925
    181 : 0.20535122
    182 : 0.14376405
    183 : 0.19113998
    184 : 0.15255031
    185 : 0.1683653
    186 : 0.1345416
    187 : 0.18488044
    188 : 0.14493594
    189 : 0.13141134
    190 : 0.12909631
    191 : 0.118659936
    192 : 0.17325518
    193 : 0.1790213
    194 : 0.1699938
    195 : 0.16877407
    196 : 0.18976636
    197 : 0.19119301
    198 : 0.18941264
    199 : 0.13560736
    200 : 0.16625994

    Saving weights

    Epoch # :  3
    201 : 0.17901282
    202 : 0.17013286
    203 : 0.17263316
    204 : 0.1718896
    205 : 0.17894654
    206 : 0.13783678
    207 : 0.1753944
    208 : 0.15459892
    209 : 0.18952
    210 : 0.1633486
    211 : 0.2168323
    212 : 0.17050663
    213 : 0.16417171
    214 : 0.16238953
    215 : 0.16125214
    216 : 0.18151295
    217 : 0.16101979
    218 : 0.1741256
    219 : 0.15554166
    220 : 0.11732925
    221 : 0.15187575
    222 : 0.14885443
    223 : 0.11471118
    224 : 0.13407499
    225 : 0.15736902
    226 : 0.21802942
    227 : 0.26981345
    228 : 0.17373471
    229 : 0.16089554
    230 : 0.16817829
    231 : 0.17552322
    232 : 0.1361039
    233 : 0.14796409
    234 : 0.13840519
    235 : 0.143079
    236 : 0.15123414
    237 : 0.14893511
    238 : 0.17936844
    239 : 0.17431177
    240 : 0.17261112
    241 : 0.18554956
    242 : 0.14208262
    243 : 0.13131864
    244 : 0.19758537
    245 : 0.14836484
    246 : 0.10673524
    247 : 0.17087989
    248 : 0.1598848
    249 : 0.19041458
    250 : 0.14542401
    251 : 0.14849691
    252 : 0.14587799
    253 : 0.15194274
    254 : 0.20370308
    255 : 0.19295849
    256 : 0.1677388
    257 : 0.18849194
    258 : 0.15776236
    259 : 0.16045012
    260 : 0.18263952
    261 : 0.17164695
    262 : 0.19055504
    263 : 0.14306639
    264 : 0.13394839
    265 : 0.112104625
    266 : 0.15312877
    267 : 0.22132248
    268 : 0.10174039
    269 : 0.2142402
    270 : 0.15739451
    271 : 0.12877306
    272 : 0.16289003
    273 : 0.09586678
    274 : 0.13516706
    275 : 0.13712475
    276 : 0.107314326
    277 : 0.13197443
    278 : 0.19999763
    279 : 0.223027
    280 : 0.1389774
    281 : 0.15779737
    282 : 0.107019894
    283 : 0.14449008
    284 : 0.18447125
    285 : 0.15488441
    286 : 0.18284425
    287 : 0.2246656
    288 : 0.1205272
    289 : 0.19679895
    290 : 0.12558122
    291 : 0.16750559
    292 : 0.17685059
    293 : 0.13895595
    294 : 0.16425319
    295 : 0.17484966
    296 : 0.17327179
    297 : 0.1464128
    298 : 0.17445475
    299 : 0.12021272
    300 : 0.11980891

    Saving weights

    Epoch # :  4
    301 : 0.16603279
    302 : 0.24273308
    303 : 0.107660726
    304 : 0.1818581
    305 : 0.15924124
    306 : 0.19342513
    307 : 0.19442216
    308 : 0.14493223
    309 : 0.17978013
    310 : 0.18150403
    311 : 0.14822552
    312 : 0.19914609
    313 : 0.22097462
    314 : 0.1526828
    315 : 0.15185863
    316 : 0.22702914
    317 : 0.13118339
    318 : 0.18537629
    319 : 0.12871832
    320 : 0.14410262
    321 : 0.18395896
    322 : 0.13001217
    323 : 0.18391185
    324 : 0.12477714
    325 : 0.199192
    326 : 0.16134998
    327 : 0.22541869
    328 : 0.16213539
    329 : 0.13522245
    330 : 0.13184895
    331 : 0.15350123
    332 : 0.15710048
    333 : 0.15941197
    334 : 0.15219384
    335 : 0.14252946
    336 : 0.1589842
    337 : 0.16354874
    338 : 0.18047932
    339 : 0.15408078
    340 : 0.17134517
    341 : 0.15374233
    342 : 0.15638477
    343 : 0.13464668
    344 : 0.13939834
    345 : 0.18057472
    346 : 0.1492266
    347 : 0.12813078
    348 : 0.17444542
    349 : 0.17342068
    350 : 0.16544814
    351 : 0.118113235
    352 : 0.15712753
    353 : 0.1633873
    354 : 0.14558057
    355 : 0.15079728
    356 : 0.13880382
    357 : 0.1671846
    358 : 0.15840809
    359 : 0.14661627
    360 : 0.11951293
    361 : 0.12784319
    362 : 0.14449261
    363 : 0.17243747
    364 : 0.14989041
    365 : 0.18783264
    366 : 0.15136482
    367 : 0.1285991
    368 : 0.17548357
    369 : 0.13375932
    370 : 0.12807311
    371 : 0.1403205
    372 : 0.16770568
    373 : 0.16714698
    374 : 0.1418988
    375 : 0.14442417
    376 : 0.15763202
    377 : 0.12764521
    378 : 0.15478568
    379 : 0.15811236
    380 : 0.13936108
    381 : 0.09586071
    382 : 0.1208115
    383 : 0.14172553
    384 : 0.12709469
    385 : 0.15568139
    386 : 0.20636272
    387 : 0.16021408
    388 : 0.11895918
    389 : 0.1862535
    390 : 0.14641564
    391 : 0.136253
    392 : 0.14310972
    393 : 0.17063469
    394 : 0.16109401
    395 : 0.18627697
    396 : 0.1678326
    397 : 0.11674715
    398 : 0.120335296
    399 : 0.16727264
    400 : 0.16988155

    Saving weights

    Epoch # :  5
    401 : 0.18953583
    402 : 0.21974212
    403 : 0.14886592
    404 : 0.17754227
    405 : 0.15967494
    406 : 0.16220465
    407 : 0.17538185
    408 : 0.13985996
    409 : 0.15688066
    410 : 0.17295417
    411 : 0.19125944
    412 : 0.15140836
    413 : 0.1515368
    414 : 0.1371538
    415 : 0.14934567
    416 : 0.16906239
    417 : 0.17437431
    418 : 0.121490926
    419 : 0.13907091
    420 : 0.1818716
    421 : 0.13308492
    422 : 0.15435827
    423 : 0.1085796
    424 : 0.17504019
    425 : 0.18501039
    426 : 0.14386661
    427 : 0.16162802
    428 : 0.12870313
    429 : 0.15728018
    430 : 0.12465912
    431 : 0.13714069
    432 : 0.15174665
    433 : 0.15864854
    434 : 0.12399879
    435 : 0.15132429
    436 : 0.16189198
    437 : 0.13215101
    438 : 0.085509084
    439 : 0.12365692
    440 : 0.14158668
    441 : 0.14126694
    442 : 0.17477725
    443 : 0.114743836
    444 : 0.14275879
    445 : 0.12334534
    446 : 0.16865003
    447 : 0.13654107
    448 : 0.14612833
    449 : 0.09927992
    450 : 0.15159115
    451 : 0.15787706
    452 : 0.13062805
    453 : 0.15452929
    454 : 0.12403348
    455 : 0.17622524
    456 : 0.12665837
    457 : 0.12294716
    458 : 0.13972677
    459 : 0.12148523
    460 : 0.15410598
    461 : 0.14123234
    462 : 0.1843355
    463 : 0.12634191
    464 : 0.12272131
    465 : 0.0899757
    466 : 0.11561899
    467 : 0.14387791
    468 : 0.14004202
    469 : 0.15491392
    470 : 0.105801255
    471 : 0.14121528
    472 : 0.17053694
    473 : 0.119935796
    474 : 0.15736677
    475 : 0.11210575
    476 : 0.10346422
    477 : 0.12853362
    478 : 0.16391693
    479 : 0.16938986
    480 : 0.11545515
    481 : 0.121684745
    482 : 0.13424827
    483 : 0.19649269
    484 : 0.13456713
    485 : 0.12803222
    486 : 0.10114731
    487 : 0.104497984
    488 : 0.14274797
    489 : 0.160999
    490 : 0.13257213
    491 : 0.1433788
    492 : 0.08464967
    493 : 0.12354956
    494 : 0.09862904
    495 : 0.1193639
    496 : 0.12384308
    497 : 0.10837276
    498 : 0.15399057
    499 : 0.14242369
    500 : 0.11715834

    Saving weights

    Epoch # :  6
    501 : 0.13442406
    502 : 0.122542895
    503 : 0.096438885
    504 : 0.11270418
    505 : 0.12368342
    506 : 0.15054221
    507 : 0.11291506
    508 : 0.08080871
    509 : 0.09131373
    510 : 0.08952438
    511 : 0.11494467
    512 : 0.09245553
    513 : 0.10379431
    514 : 0.074560665
    515 : 0.07031863
    516 : 0.079880826
    517 : 0.091697544
    518 : 0.08455921
    519 : 0.07518372
    520 : 0.093540765
    521 : 0.073418446
    522 : 0.05932852
    523 : 0.08722296
    524 : 0.084686145
    525 : 0.068582505
    526 : 0.053134985
    527 : 0.08068697
    528 : 0.06336024
    529 : 0.048133183
    530 : 0.06521394
    531 : 0.0921517
    532 : 0.066156454
    533 : 0.070156984
    534 : 0.06810923
    535 : 0.07680989
    536 : 0.05213432
    537 : 0.06263829
    538 : 0.06623813
    539 : 0.05790771
    540 : 0.06402144
    541 : 0.065428525
    542 : 0.0704135
    543 : 0.057756376
    544 : 0.08139511
    545 : 0.07275272
    546 : 0.048887666
    547 : 0.08097727
    548 : 0.060200233
    549 : 0.053648386
    550 : 0.062323753
    551 : 0.056413125
    552 : 0.058201514
    553 : 0.05253775
    554 : 0.054922543
    555 : 0.063353494
    556 : 0.054380465
    557 : 0.0602294
    558 : 0.053169355
    559 : 0.053491496
    560 : 0.042410105
    561 : 0.07057136
    562 : 0.055632524
    563 : 0.067340404
    564 : 0.08109722
    565 : 0.034673423
    566 : 0.060106255
    567 : 0.048393518
    568 : 0.053593513
    569 : 0.06268452
    570 : 0.045645233
    571 : 0.055787552
    572 : 0.0441093
    573 : 0.044084158
    574 : 0.03573328
    575 : 0.052873623
    576 : 0.04968297
    577 : 0.039845023
    578 : 0.03784553
    579 : 0.03355617
    580 : 0.04751711
    581 : 0.04422307
    582 : 0.05189165
    583 : 0.04482298
    584 : 0.04841815
    585 : 0.04999983
    586 : 0.037030555
    587 : 0.03569565
    588 : 0.087424584
    589 : 0.04230785
    590 : 0.08994114
    591 : 0.059833102
    592 : 0.0658737
    593 : 0.07329171
    594 : 0.04695451
    595 : 0.05152605
    596 : 0.08095735
    597 : 0.051839095
    598 : 0.055347454
    599 : 0.053730503
    600 : 0.036761034

    Saving weights

    Epoch # :  7
    601 : 0.04426942
    602 : 0.043159552
    603 : 0.04522834
    604 : 0.039966177
    605 : 0.035672095
    606 : 0.042196684
    607 : 0.034805223
    608 : 0.03537173
    609 : 0.040783793
    610 : 0.03544075
    611 : 0.03690445
    612 : 0.049550585
    613 : 0.033476066
    614 : 0.05798296
    615 : 0.054652005
    616 : 0.05621876
    617 : 0.04641015
    618 : 0.03709464
    619 : 0.0301331
    620 : 0.022059407
    621 : 0.034597322
    622 : 0.031191468
    623 : 0.037477817
    624 : 0.034152936
    625 : 0.027973862
    626 : 0.031034833
    627 : 0.03919048
    628 : 0.034658086
    629 : 0.031886067
    630 : 0.033454116
    631 : 0.029404182
    632 : 0.019120205
    633 : 0.028540902
    634 : 0.027149694
    635 : 0.025564961
    636 : 0.03381304
    637 : 0.03865838
    638 : 0.023045046
    639 : 0.027993126
    640 : 0.025458282
    641 : 0.017809348
    642 : 0.023424316
    643 : 0.021662088
    644 : 0.02771093
    645 : 0.035130486
    646 : 0.018692678
    647 : 0.020653889
    648 : 0.02913595
    649 : 0.01891705
    650 : 0.017994203
    651 : 0.023228765
    652 : 0.021570383
    653 : 0.02654841
    654 : 0.025029993
    655 : 0.027426854
    656 : 0.05428234
    657 : 0.020585785
    658 : 0.030386603
    659 : 0.022920083
    660 : 0.055720307
    661 : 0.022931974
    662 : 0.04760602
    663 : 0.021576257
    664 : 0.046356615
    665 : 0.030228818
    666 : 0.028182551
    667 : 0.0395816
    668 : 0.03088934
    669 : 0.03424506
    670 : 0.03019856
    671 : 0.02267797
    672 : 0.04645315
    673 : 0.030178037
    674 : 0.016843896
    675 : 0.04730647
    676 : 0.02467842
    677 : 0.024688885
    678 : 0.028593015
    679 : 0.026307387
    680 : 0.026936525
    681 : 0.033172917
    682 : 0.024347525
    683 : 0.02469768
    684 : 0.039411154
    685 : 0.021450425
    686 : 0.03696449
    687 : 0.02171352
    688 : 0.027282929
    689 : 0.031916257
    690 : 0.02235554
    691 : 0.017959505
    692 : 0.020838968
    693 : 0.013252987
    694 : 0.021836024
    695 : 0.018658968
    696 : 0.013759667
    697 : 0.023081947
    698 : 0.013124827
    699 : 0.017545635
    700 : 0.020236535

    Saving weights

    Epoch # :  8
    701 : 0.018621003
    702 : 0.019353526
    703 : 0.015889335
    704 : 0.022002682
    705 : 0.018324777
    706 : 0.018965067
    707 : 0.012721606
    708 : 0.019445775
    709 : 0.014131368
    710 : 0.01161105
    711 : 0.015342333
    712 : 0.018221237
    713 : 0.015236407
    714 : 0.020059306
    715 : 0.013945323
    716 : 0.01167308
    717 : 0.01594623
    718 : 0.021121645
    719 : 0.00896197
    720 : 0.036729336
    721 : 0.021395124
    722 : 0.02969343
    723 : 0.019755252
    724 : 0.021339756
    725 : 0.023190228
    726 : 0.02650289
    727 : 0.02220418
    728 : 0.019846085
    729 : 0.014317309
    730 : 0.022206632
    731 : 0.021244489
    732 : 0.019820249
    733 : 0.025693035
    734 : 0.01064799
    735 : 0.015685325
    736 : 0.019062035
    737 : 0.012731779
    738 : 0.016261913
    739 : 0.021456754
    740 : 0.011810435
    741 : 0.033399843
    742 : 0.020201964
    743 : 0.025308514
    744 : 0.023638744
    745 : 0.010511067
    746 : 0.02432734
    747 : 0.018453607
    748 : 0.020367747
    749 : 0.021423418
    750 : 0.018755998
    751 : 0.018316936
    752 : 0.016256455
    753 : 0.014962328
    754 : 0.015892908
    755 : 0.014143806
    756 : 0.012846584
    757 : 0.0177399
    758 : 0.015199447
    759 : 0.014984873
    760 : 0.009192971
    761 : 0.014841373
    762 : 0.0137243625
    763 : 0.013109165
    764 : 0.0077580977
    765 : 0.01603927
    766 : 0.019640386
    767 : 0.013697833
    768 : 0.033287555
    769 : 0.029298002
    770 : 0.021584593
    771 : 0.043512568
    772 : 0.014624441
    773 : 0.029810317
    774 : 0.035879403
    775 : 0.014115298
    776 : 0.016951647
    777 : 0.03162256
    778 : 0.011991432
    779 : 0.016980309
    780 : 0.035782706
    781 : 0.021677561
    782 : 0.021091742
    783 : 0.029512387
    784 : 0.02647265
    785 : 0.024030041
    786 : 0.024929836
    787 : 0.018063005
    788 : 0.008785535
    789 : 0.020103712
    790 : 0.015540041
    791 : 0.01722282
    792 : 0.021800814
    793 : 0.02077612
    794 : 0.011893157
    795 : 0.03442281
    796 : 0.018371
    797 : 0.017826108
    798 : 0.020176275
    799 : 0.0132283345
    800 : 0.020353246

    Saving weights

    Epoch # :  9
    801 : 0.022759078
    802 : 0.013114662
    803 : 0.017137632
    804 : 0.02172373
    805 : 0.016473189
    806 : 0.01178815
    807 : 0.013460518
    808 : 0.023368368
    809 : 0.020377303
    810 : 0.020653604
    811 : 0.015075467
    812 : 0.015181107
    813 : 0.019734774
    814 : 0.016633855
    815 : 0.013387004
    816 : 0.024626108
    817 : 0.00936027
    818 : 0.013193708
    819 : 0.0207255
    820 : 0.009926654
    821 : 0.01382466
    822 : 0.017765138
    823 : 0.012735789
    824 : 0.014075017
    825 : 0.019548077
    826 : 0.015959276
    827 : 0.014087203
    828 : 0.013689407
    829 : 0.013260215
    830 : 0.014903953
    831 : 0.015984047
    832 : 0.017570024
    833 : 0.01403746
    834 : 0.017031007
    835 : 0.012779298
    836 : 0.015140491
    837 : 0.011804497
    838 : 0.014611011
    839 : 0.015867423
    840 : 0.012917713
    841 : 0.0141227655
    842 : 0.015325542
    843 : 0.01193547
    844 : 0.014120707
    845 : 0.010931286
    846 : 0.017071448
    847 : 0.009339923
    848 : 0.011826693
    849 : 0.0136071285
    850 : 0.014201772
    851 : 0.011183036
    852 : 0.015801456
    853 : 0.013410414
    854 : 0.019858072
    855 : 0.010668758
    856 : 0.011055729
    857 : 0.009397471
    858 : 0.011236489
    859 : 0.010794222
    860 : 0.011504445
    861 : 0.014558461
    862 : 0.013604638
    863 : 0.01736072
    864 : 0.017214041
    865 : 0.010631552
    866 : 0.018480306
    867 : 0.01355107
    868 : 0.014801864
    869 : 0.016677847
    870 : 0.01124115
    871 : 0.010724196
    872 : 0.019378453
    873 : 0.012815522
    874 : 0.013748726
    875 : 0.015646065
    876 : 0.009978996
    877 : 0.017821662
    878 : 0.018901743
    879 : 0.011759674
    880 : 0.012253387
    881 : 0.011928685
    882 : 0.014324472
    883 : 0.016857965
    884 : 0.024497246
    885 : 0.015213282
    886 : 0.0137779685
    887 : 0.015823575
    888 : 0.014339625
    889 : 0.0113422405
    890 : 0.015732324
    891 : 0.011432502
    892 : 0.009903278
    893 : 0.021576665
    894 : 0.016072292
    895 : 0.013661051
    896 : 0.011698194
    897 : 0.014100059
    898 : 0.015998408
    899 : 0.01366734
    900 : 0.017228173

    Saving weights

    Epoch # :  10
    901 : 0.015256517
    902 : 0.017925426
    903 : 0.018830191
    904 : 0.010235155
    905 : 0.01824233
    906 : 0.012391191
    907 : 0.008962083
    908 : 0.010489719
    909 : 0.014917526
    910 : 0.009348296
    911 : 0.015862988
    912 : 0.0158716
    913 : 0.0123006655
    914 : 0.010947484
    915 : 0.011996503
    916 : 0.009331629
    917 : 0.01412625
    918 : 0.011858358
    919 : 0.015962321
    920 : 0.015260791
    921 : 0.012145869
    922 : 0.012631931
    923 : 0.012133024
    924 : 0.01436346
    925 : 0.0090621915
    926 : 0.018338593
    927 : 0.013865862
    928 : 0.015121463
    929 : 0.012715915
    930 : 0.015406801
    931 : 0.007944118
    932 : 0.013012178
    933 : 0.015375209
    934 : 0.011021354
    935 : 0.011736629
    936 : 0.013288257
    937 : 0.01125223
    938 : 0.01520826
    939 : 0.009739559
    940 : 0.0095591955
    941 : 0.012369012
    942 : 0.013733639
    943 : 0.015813028
    944 : 0.014613378
    945 : 0.0108814025
    946 : 0.017073147
    947 : 0.012277768
    948 : 0.011523608
    949 : 0.010260019
    950 : 0.0105614625
    951 : 0.010187719
    952 : 0.012835086
    953 : 0.008723493
    954 : 0.011253092
    955 : 0.0161443
    956 : 0.0138244815
    957 : 0.016399836
    958 : 0.012679138
    959 : 0.010885792
    960 : 0.015541065
    961 : 0.021089569
    962 : 0.009084817
    963 : 0.015007372
    964 : 0.017662758
    965 : 0.010010967
    966 : 0.00870137
    967 : 0.014901962
    968 : 0.009409742
    969 : 0.01382525
    970 : 0.012637113
    971 : 0.022212887
    972 : 0.012161435
    973 : 0.025501728
    974 : 0.010741842
    975 : 0.010838392
    976 : 0.012944105
    977 : 0.01013847
    978 : 0.011446202
    979 : 0.0132047795
    980 : 0.015261328
    981 : 0.010831416
    982 : 0.009888933
    983 : 0.012245169
    984 : 0.009549674
    985 : 0.008401048
    986 : 0.011405158
    987 : 0.009189022
    988 : 0.012930982
    989 : 0.0069505097
    990 : 0.015047225
    991 : 0.009100561
    992 : 0.00910202
    993 : 0.0137563525
    994 : 0.008319947
    995 : 0.010955427
    996 : 0.008110689
    997 : 0.01019202
    998 : 0.008356452
    999 : 0.00993902
    1000 : 0.010270872

    Saving weights

```python
import matplotlib.pyplot as plt

# remove noise from initial epochs
loss_history_plot = list(filter(lambda x: x < 0.25, loss_history))

plt.figure(figsize=(12, 5))
plt.plot(loss_history_plot)
plt.show()
```

# Observations
The model does train for 100 timesteps quite quickly, but it doesnt match the speed of learning by the paper. In the paper, the loss starts to drop close to the 300th iteration and drops all the way to the low 0.0x by the 425th iteration.

Here the loss starts to reduce at the 400th iteration and drops completely by the 600th iteration, much slower than the paper. Also, I had to clip off the first few iterations with extremely high losses as the graph cannot be visualized properly if they were kept.

#  Loading Checkpoints

Tensorflow Eager makes loading of checkpoints that are fully contained in a Keras model easy. However, there is one small idiosynchacy that needs to be dealt with when loading models in TF 1.7.

Models, after being build, need to be called at least once prior to the checkpoint restoring their weights. If this is not done, it throws obscure errors.

```python
if os.path.exists(CHECKPOINTS_DIR):
    ckpt_path = tf.train.latest_checkpoint(CHECKPOINTS_DIR)
    print("Checkpoint path : ", ckpt_path)

    with tf.device(_device):
        model = JANETModel(NUM_UNITS, num_outputs=1, num_timesteps=TIME_STEPS, output_activation='linear')

        # model.compile(tf.train.AdamOptimizer(), loss='mse')

        # this is where you need to call the model at least once,
        # so that all of its variables can be properly restored
        zeros = tf.zeros((1, TIME_STEPS, 2))
        model(zeros)

        # restore the weights
        tfe.Checkpoint(model=model, ).restore(ckpt_path)

        # predict one batch to ensure the weights are correctly loaded
        generator = batch_generator()

        losses = []
        for i in range(20):
            inputs, outputs = next(generator)

            preds = model(inputs)
            loss_val = loss(outputs, preds)

            losses.append(loss_val.numpy())

    print("Final average predicted error (should be less than 0.03) : ", np.mean(losses))
```

    Checkpoint path :  checkpoints_addition/addition_janet.ckpt-1
    Final average predicted error (should be less than 0.03) :  0.010359424

# Stepping off the well trodden path
When building keras models that **only** use Keras layers, there is no issue when loading or saving weights.

However, not every thing can be represented in layers, and a few models may require working directlty with TF Eager variables. In such cases, a different approach is required to save and restore these models.

Lets see the same model without using Keras layers to define them. Here to keep it simple, we will use Eager variables to define the RNN cell, and a Keras Layer to define the classification / regression layer.

```python
class EagerJANETModel(tf.keras.Model):

    def __init__(self, input_dim, units, num_outputs, num_timesteps, output_activation='sigmoid', **kwargs):
        super(EagerJANETModel, self).__init__(**kwargs)

        self.input_dim = input_dim
        self.units = units
        self.classes = num_outputs
        self.num_timesteps = num_timesteps

        # Initialize the forget gate with ChronoInitializer
        # The memory gate is initialized with zeros
        def bias_initializer(_, *args, **kwargs):
            forget_gate = ChronoInitializer(self.num_timesteps)((self.units,), *args, **kwargs)

            return tf.keras.backend.concatenate([
                forget_gate,
                tf.keras.initializers.Zeros()((self.units,), *args, **kwargs),
            ])

        # Initialize the hidden and recurrent gates
        self.kernel = tf.get_variable('kernel', shape=[input_dim, units * 2], dtype=tf.float32,
                                     initializer=tf.keras.initializers.glorot_uniform())

        self.recurrent_kernel = tf.get_variable('recurrent_kernel', shape=[units, units * 2], dtype=tf.float32,
                                                initializer=tf.keras.initializers.glorot_uniform())

        self.recurrent_bias = tf.get_variable('recurrent_bias', shape=[units * 2], dtype=tf.float32,
                                              initializer=bias_initializer)

        # Initialize the final layer (for classification or regression depending on the output activation)
        self.output_dense = tf.keras.layers.Dense(num_outputs, activation=output_activation)

        # we need to create a dictionary of all of the weights which are not in Keras layers
        self.additional_weights = OrderedDict()
        self.additional_weights[self.kernel.name] = self.kernel
        self.additional_weights[self.recurrent_kernel.name] = self.recurrent_kernel
        self.additional_weights[self.recurrent_bias.name] = self.recurrent_bias

    """
    This is required to override a certain issue with how the weights of a model are loaded
    by a checkpoint.

    While we could use `model.call(inputs)` directly, the more pythonic way of doing this is
    by using `model(inputs)`.

    However, in Eager mode, the ordinary function call does *not* forward the method to the
    models `call` method. This bypasses that issue.
    """
    def __call__(self, *args, **kwargs):
        if not tf.executing_eagerly():
            super(EagerJANETModel, self).__call__(*args, **kwargs)
        return self.call(*args, **kwargs)

    def call(self, inputs, training=None, mask=None):
        # Initialize the hidden memory states
        outputs = []
        states = []

        h_state = tf.zeros((inputs.shape[0], self.units))
        c_state = tf.zeros((inputs.shape[0], self.units))

        # Input is in the shape [None, timesteps, input_dim]
        for t in range(inputs.shape[1]):
            ip = inputs[:, t, :]  # access the T'th timestep

            # Perform the forward pass of the model
            z = tf.matmul(ip, self.kernel)
            z += tf.matmul(h_state, self.recurrent_kernel) + self.recurrent_bias

            # Split the output into the forget and memory outputs
            z0 = z[:, :self.units]
            z1 = z[:, self.units: 2 * self.units]

            # gate updates
            f = tf.keras.activations.sigmoid(z0)
            c = f * c_state + (1. - f) * tf.nn.tanh(z1)

            # state updates
            h = c

            # update our previous state
            h_state = h
            c_state = c

            # preserve the history of our states
            outputs.append(h)
            states.append([h])  # here we maintain just 1 state variable, can be more

        # save memory states to global variables so we can access them later
        # self.cell_outputs = tf.stack(outputs, axis=1)
        # self.cell_states = states

        # perform forward pass of the classifier / regressor
        preds = self.output_dense(outputs[-1])

        return preds
```

# Gradients of a custom model

The gradients of such a model are a little more complicated than the earlier simple `model.trainable_variables`. Since tf.keras Models do not track the weights that are added separately (apart from keras layers, no weights or variables are managed by Keras), we need to write a custom grad function.

```python
def grad(model, X, y):
    with tf.GradientTape() as tape:
        preds = model(X)
        loss_val = loss(y, preds)

    # this is the crucial step : use the dictionary of weights that we manage
    # manually to get all of the weights that are not managed by keras and add
    # them to the list of weights for which we need gradients
    vars = model.trainable_variables + list(model.additional_weights.values())
    grads = tape.gradient(loss_val, vars)
    grad_vars = zip(grads, vars)

    return grad_vars, loss_val
```

# Reason for custom gradient function

Keras has two very convenient methods `add_weight` and `add_variable` to add weights to a layer / model easily. However, TF Eager mode has not implemented these methods yet, and therefore we cant use them directly.

Therefore, we have to take the roundabout way of mixing Keras layers with custom weights and keep a track of them using a dictionary.

# Training and Saving the weights of a custom model

Saving and restoring of custom models is exactly the same as before if you are using tf.get_variable to create variables in Eager mode. However, tfe.Variable() is NOT managed by keras, so don't use it inside a Keras model.

```python
%%time
input_dim = 2

with tf.device(_device):
    model = EagerJANETModel(input_dim, NUM_UNITS, num_outputs=1, num_timesteps=TIME_STEPS, output_activation='linear')

    optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE)
    global_step = tf.train.get_or_create_global_step()

    # try using different optimizers and different optimizer configs
    # model.compile(loss='mse', optimizer=optimizer)

    best_loss = 100.
    generator = batch_generator()

loss_history = []
update_counter = 1

for epoch in range(NUM_EPOCHS):
    print("Epoch # : ", epoch + 1)

    for step in range(STEPS_PER_EPOCH):
        # get batch dataset
        with tf.device(_device):
            inputs, targets = next(generator)

            # get gradients and loss at this iteration (uses the new grad method)
            gradients, loss_val = grad(model, inputs, targets)

            # apply gradients
            optimizer.apply_gradients(gradients, tf.train.get_or_create_global_step())

        loss_history.append(loss_val.numpy())
        print(update_counter, ":", loss_history[-1])

        update_counter += 1

    print()

    # remove old checkpoint that we no longer need
    if os.path.exists(CHECKPOINTS_DIR):
        shutil.rmtree(CHECKPOINTS_DIR)

    # save the checkpoint weights
    # Note how we add the model's additional weights as name-variable pairs
    checkpoint = tfe.Checkpoint(model=model).save(CHECKPOINT_PATH)

    # Optional : Save the weight matrices in Keras format as well
    #model.save_weights('addition_model.h5', overwrite=True)

    print("Saving weights")
    print()
print()

```

    Epoch # :  1
    1 : 1.1070119
    2 : 1.1105163
    3 : 0.8818216
    4 : 1.0917069
    5 : 0.77383345
    6 : 1.0805733
    7 : 0.83379626
    8 : 0.7731505
    9 : 0.55083674
    10 : 0.54209757
    11 : 0.28447378
    12 : 0.22762008
    13 : 0.13934982
    14 : 0.20197843
    15 : 0.39939064
    16 : 0.31061924
    17 : 0.15716037
    18 : 0.18323703
    19 : 0.13392271
    20 : 0.13464299
    21 : 0.276382
    22 : 0.17964411
    23 : 0.22904907
    24 : 0.21105047
    25 : 0.26379216
    26 : 0.20097953
    27 : 0.1950277
    28 : 0.21655178
    29 : 0.26505485
    30 : 0.20399876
    31 : 0.17650402
    32 : 0.16045463
    33 : 0.17143913
    34 : 0.14122799
    35 : 0.20478329
    36 : 0.18234424
    37 : 0.23561493
    38 : 0.17591843
    39 : 0.17281689
    40 : 0.15329581
    41 : 0.16284627
    42 : 0.17922379
    43 : 0.10488374
    44 : 0.19159271
    45 : 0.18599848
    46 : 0.13775517
    47 : 0.20575161
    48 : 0.2778363
    49 : 0.18510982
    50 : 0.16983838
    51 : 0.19744366
    52 : 0.13969105
    53 : 0.18370225
    54 : 0.14061068
    55 : 0.1928698
    56 : 0.17589724
    57 : 0.2265582
    58 : 0.1547241
    59 : 0.12760934
    60 : 0.19300795
    61 : 0.18235433
    62 : 0.1795494
    63 : 0.17900254
    64 : 0.21244241
    65 : 0.20464407
    66 : 0.15752582
    67 : 0.15200463
    68 : 0.14627606
    69 : 0.18742035
    70 : 0.16700277
    71 : 0.1247301
    72 : 0.16766438
    73 : 0.14763501
    74 : 0.20850983
    75 : 0.19158737
    76 : 0.20135303
    77 : 0.1800587
    78 : 0.22546035
    79 : 0.140615
    80 : 0.1487796
    81 : 0.14596748
    82 : 0.18284404
    83 : 0.14875966
    84 : 0.13198483
    85 : 0.1530276
    86 : 0.24958718
    87 : 0.19805336
    88 : 0.14327373
    89 : 0.15821955
    90 : 0.17935774
    91 : 0.12539722
    92 : 0.12265874
    93 : 0.2019233
    94 : 0.15431829
    95 : 0.14339465
    96 : 0.12539534
    97 : 0.13147862
    98 : 0.16164501
    99 : 0.14825484
    100 : 0.12537865

    Saving weights

    Epoch # :  2
    101 : 0.18204308
    102 : 0.17038606
    103 : 0.17528112
    104 : 0.14784223
    105 : 0.1736657
    106 : 0.19151993
    107 : 0.12320281
    108 : 0.150907
    109 : 0.11942315
    110 : 0.1559823
    111 : 0.13425086
    112 : 0.18812114
    113 : 0.17583229
    114 : 0.16656788
    115 : 0.17426282
    116 : 0.21648937
    117 : 0.14646952
    118 : 0.16117752
    119 : 0.15734364
    120 : 0.13107187
    121 : 0.13035534
    122 : 0.18611099
    123 : 0.14148153
    124 : 0.15214053
    125 : 0.13172022
    126 : 0.17997101
    127 : 0.17663139
    128 : 0.12589602
    129 : 0.16294645
    130 : 0.1730654
    131 : 0.1273704
    132 : 0.1717456
    133 : 0.1745816
    134 : 0.17692728
    135 : 0.16464111
    136 : 0.18314973
    137 : 0.13289616
    138 : 0.13820094
    139 : 0.17133106
    140 : 0.15152824
    141 : 0.17567623
    142 : 0.17623359
    143 : 0.19180441
    144 : 0.18352692
    145 : 0.13728999
    146 : 0.17852764
    147 : 0.12572375
    148 : 0.16284056
    149 : 0.1642189
    150 : 0.17429633
    151 : 0.1435761
    152 : 0.11713633
    153 : 0.14107157
    154 : 0.14546867
    155 : 0.1790529
    156 : 0.17457055
    157 : 0.17983647
    158 : 0.14202647
    159 : 0.15026477
    160 : 0.15887919
    161 : 0.12813525
    162 : 0.17607982
    163 : 0.16810997
    164 : 0.13685669
    165 : 0.12371718
    166 : 0.16965851
    167 : 0.1762833
    168 : 0.1958768
    169 : 0.17784019
    170 : 0.17873158
    171 : 0.15622854
    172 : 0.109150365
    173 : 0.18188292
    174 : 0.13565944
    175 : 0.117799304
    176 : 0.20079678
    177 : 0.16492237
    178 : 0.1638561
    179 : 0.16425999
    180 : 0.20874603
    181 : 0.20228085
    182 : 0.15189184
    183 : 0.18769886
    184 : 0.17753197
    185 : 0.18713894
    186 : 0.14705884
    187 : 0.15809944
    188 : 0.123830624
    189 : 0.17105064
    190 : 0.121779546
    191 : 0.12211758
    192 : 0.18868004
    193 : 0.20019177
    194 : 0.15259933
    195 : 0.1594625
    196 : 0.1961345
    197 : 0.15698275
    198 : 0.17575195
    199 : 0.11148106
    200 : 0.20747012

    Saving weights

    Epoch # :  3
    201 : 0.14887035
    202 : 0.19070725
    203 : 0.1713575
    204 : 0.17270532
    205 : 0.13449092
    206 : 0.14817455
    207 : 0.13158591
    208 : 0.12642026
    209 : 0.14794159
    210 : 0.18951969
    211 : 0.15350872
    212 : 0.17310476
    213 : 0.18515782
    214 : 0.15196788
    215 : 0.16284847
    216 : 0.1526328
    217 : 0.14186777
    218 : 0.16717991
    219 : 0.14211625
    220 : 0.15516815
    221 : 0.14800154
    222 : 0.16182373
    223 : 0.21034464
    224 : 0.11586529
    225 : 0.17499305
    226 : 0.20132546
    227 : 0.19474964
    228 : 0.16007635
    229 : 0.1592393
    230 : 0.16491275
    231 : 0.18530732
    232 : 0.14145273
    233 : 0.1768079
    234 : 0.13191521
    235 : 0.18317261
    236 : 0.15648174
    237 : 0.1877356
    238 : 0.15426111
    239 : 0.16922584
    240 : 0.18376307
    241 : 0.11959646
    242 : 0.14797199
    243 : 0.17166832
    244 : 0.14980684
    245 : 0.18748964
    246 : 0.14002205
    247 : 0.13780145
    248 : 0.14275318
    249 : 0.14382993
    250 : 0.1838832
    251 : 0.14226799
    252 : 0.14741318
    253 : 0.19078712
    254 : 0.15610561
    255 : 0.17482343
    256 : 0.1589665
    257 : 0.1183536
    258 : 0.18757202
    259 : 0.16272826
    260 : 0.13007227
    261 : 0.15286765
    262 : 0.15288132
    263 : 0.13993363
    264 : 0.17197958
    265 : 0.18528095
    266 : 0.16949825
    267 : 0.12199705
    268 : 0.21682921
    269 : 0.15110467
    270 : 0.21232148
    271 : 0.22917467
    272 : 0.14665098
    273 : 0.18296021
    274 : 0.15262493
    275 : 0.14605796
    276 : 0.1268351
    277 : 0.16669087
    278 : 0.11745579
    279 : 0.18110761
    280 : 0.16277245
    281 : 0.2268802
    282 : 0.20069034
    283 : 0.13055496
    284 : 0.11070875
    285 : 0.17924136
    286 : 0.14838962
    287 : 0.17903607
    288 : 0.18603896
    289 : 0.1696229
    290 : 0.14411019
    291 : 0.1300347
    292 : 0.13376462
    293 : 0.21308708
    294 : 0.12752213
    295 : 0.121160544
    296 : 0.21067172
    297 : 0.15202391
    298 : 0.13051125
    299 : 0.110127084
    300 : 0.12955508

    Saving weights

    Epoch # :  4
    301 : 0.122353114
    302 : 0.18318777
    303 : 0.13638273
    304 : 0.12237187
    305 : 0.123392105
    306 : 0.15292339
    307 : 0.12859733
    308 : 0.16218194
    309 : 0.19448827
    310 : 0.18691643
    311 : 0.15002903
    312 : 0.17923756
    313 : 0.19412674
    314 : 0.19068882
    315 : 0.144021
    316 : 0.15869999
    317 : 0.106259115
    318 : 0.19035427
    319 : 0.19068772
    320 : 0.1505567
    321 : 0.11605667
    322 : 0.109950714
    323 : 0.13136855
    324 : 0.13946646
    325 : 0.11518186
    326 : 0.19323696
    327 : 0.17520142
    328 : 0.13721709
    329 : 0.23227428
    330 : 0.16347532
    331 : 0.11325505
    332 : 0.13729224
    333 : 0.1468662
    334 : 0.19002213
    335 : 0.12993678
    336 : 0.1820619
    337 : 0.1907936
    338 : 0.175592
    339 : 0.12050202
    340 : 0.13999718
    341 : 0.16322201
    342 : 0.14092475
    343 : 0.14185663
    344 : 0.1449196
    345 : 0.16047326
    346 : 0.12820634
    347 : 0.1000732
    348 : 0.1751894
    349 : 0.18094082
    350 : 0.112328745
    351 : 0.14751756
    352 : 0.14207111
    353 : 0.11712582
    354 : 0.16550179
    355 : 0.16431375
    356 : 0.1956683
    357 : 0.14979157
    358 : 0.16625123
    359 : 0.13688497
    360 : 0.19555968
    361 : 0.13767569
    362 : 0.13543336
    363 : 0.14541203
    364 : 0.09792575
    365 : 0.17996818
    366 : 0.10850013
    367 : 0.17931066
    368 : 0.09853659
    369 : 0.21194105
    370 : 0.14727741
    371 : 0.13250084
    372 : 0.13957658
    373 : 0.17154188
    374 : 0.122357346
    375 : 0.13729267
    376 : 0.13760594
    377 : 0.14450608
    378 : 0.13199902
    379 : 0.13282973
    380 : 0.13617204
    381 : 0.13682877
    382 : 0.15752833
    383 : 0.13570952
    384 : 0.14470096
    385 : 0.13399546
    386 : 0.23874325
    387 : 0.12547314
    388 : 0.17041811
    389 : 0.1828761
    390 : 0.1721962
    391 : 0.15522917
    392 : 0.13532679
    393 : 0.18645695
    394 : 0.10208627
    395 : 0.13191958
    396 : 0.18120922
    397 : 0.13529138
    398 : 0.10852894
    399 : 0.12240328
    400 : 0.1628542

    Saving weights

    Epoch # :  5
    401 : 0.14419122
    402 : 0.14095324
    403 : 0.13627341
    404 : 0.19891028
    405 : 0.16044183
    406 : 0.14136936
    407 : 0.12900198
    408 : 0.12038742
    409 : 0.13362068
    410 : 0.12893705
    411 : 0.16242217
    412 : 0.12153445
    413 : 0.14050128
    414 : 0.089259736
    415 : 0.12443649
    416 : 0.121129096
    417 : 0.13468699
    418 : 0.09146298
    419 : 0.158473
    420 : 0.16767937
    421 : 0.12171175
    422 : 0.14509767
    423 : 0.14541373
    424 : 0.11587059
    425 : 0.13824758
    426 : 0.18851082
    427 : 0.12517287
    428 : 0.16741215
    429 : 0.19478764
    430 : 0.11757394
    431 : 0.14925705
    432 : 0.16142291
    433 : 0.15468661
    434 : 0.11863825
    435 : 0.122989714
    436 : 0.08883213
    437 : 0.11408632
    438 : 0.10624413
    439 : 0.0877149
    440 : 0.1259022
    441 : 0.14388075
    442 : 0.1629444
    443 : 0.16048867
    444 : 0.1092285
    445 : 0.16356468
    446 : 0.100446604
    447 : 0.12762791
    448 : 0.098228954
    449 : 0.14838654
    450 : 0.14381258
    451 : 0.08555947
    452 : 0.103533134
    453 : 0.14472194
    454 : 0.11253196
    455 : 0.103589
    456 : 0.14038797
    457 : 0.11605499
    458 : 0.0809207
    459 : 0.07124468
    460 : 0.11968736
    461 : 0.10753468
    462 : 0.14137371
    463 : 0.13993138
    464 : 0.10848882
    465 : 0.076820314
    466 : 0.09645891
    467 : 0.099323064
    468 : 0.081891924
    469 : 0.104056455
    470 : 0.1132332
    471 : 0.1059945
    472 : 0.124769114
    473 : 0.07393903
    474 : 0.060621336
    475 : 0.07922053
    476 : 0.06896666
    477 : 0.12955132
    478 : 0.07282942
    479 : 0.07480906
    480 : 0.0748785
    481 : 0.060773153
    482 : 0.084540784
    483 : 0.06373225
    484 : 0.10494666
    485 : 0.06685897
    486 : 0.14468968
    487 : 0.083180495
    488 : 0.06987886
    489 : 0.097393304
    490 : 0.073100865
    491 : 0.11486647
    492 : 0.09194204
    493 : 0.0654128
    494 : 0.08316186
    495 : 0.11223145
    496 : 0.07783792
    497 : 0.12735291
    498 : 0.06786843
    499 : 0.10243812
    500 : 0.06703848

    Saving weights

    Epoch # :  6
    501 : 0.07971802
    502 : 0.064800866
    503 : 0.11394337
    504 : 0.065818876
    505 : 0.0714549
    506 : 0.08406609
    507 : 0.05725936
    508 : 0.07875797
    509 : 0.061405696
    510 : 0.063146144
    511 : 0.07170342
    512 : 0.050907765
    513 : 0.054377183
    514 : 0.059061203
    515 : 0.05365225
    516 : 0.07325519
    517 : 0.064301625
    518 : 0.07002276
    519 : 0.09850267
    520 : 0.05637608
    521 : 0.116176955
    522 : 0.070346825
    523 : 0.0635022
    524 : 0.10820709
    525 : 0.077265136
    526 : 0.046696614
    527 : 0.082699515
    528 : 0.07663349
    529 : 0.05269104
    530 : 0.06779609
    531 : 0.07755826
    532 : 0.06394911
    533 : 0.08462149
    534 : 0.061114777
    535 : 0.049828798
    536 : 0.06738005
    537 : 0.082841374
    538 : 0.07510481
    539 : 0.06940232
    540 : 0.06893617
    541 : 0.08840316
    542 : 0.048258692
    543 : 0.08659206
    544 : 0.06734221
    545 : 0.07137933
    546 : 0.088551424
    547 : 0.09478065
    548 : 0.05735923
    549 : 0.06288223
    550 : 0.08321084
    551 : 0.09369582
    552 : 0.047643326
    553 : 0.064082965
    554 : 0.09248416
    555 : 0.053442735
    556 : 0.06335781
    557 : 0.094173566
    558 : 0.061406747
    559 : 0.06568434
    560 : 0.05918522
    561 : 0.061798356
    562 : 0.07224805
    563 : 0.07897055
    564 : 0.050742224
    565 : 0.072827436
    566 : 0.07351509
    567 : 0.07590517
    568 : 0.07009921
    569 : 0.042571355
    570 : 0.034844898
    571 : 0.069173954
    572 : 0.050696388
    573 : 0.066924326
    574 : 0.05468992
    575 : 0.04726695
    576 : 0.04467712
    577 : 0.064267114
    578 : 0.06235404
    579 : 0.05613732
    580 : 0.04992545
    581 : 0.057989217
    582 : 0.0393674
    583 : 0.04610258
    584 : 0.047854327
    585 : 0.046328038
    586 : 0.05036417
    587 : 0.04769492
    588 : 0.071276315
    589 : 0.05272339
    590 : 0.048787598
    591 : 0.05957798
    592 : 0.048448753
    593 : 0.059000034
    594 : 0.055239182
    595 : 0.040203203
    596 : 0.073817395
    597 : 0.04299683
    598 : 0.061882082
    599 : 0.035685763
    600 : 0.046795744

    Saving weights

    Epoch # :  7
    601 : 0.048181295
    602 : 0.051347405
    603 : 0.060216364
    604 : 0.045976263
    605 : 0.039950233
    606 : 0.0330235
    607 : 0.05420722
    608 : 0.03540788
    609 : 0.03610782
    610 : 0.046176564
    611 : 0.035257515
    612 : 0.034568172
    613 : 0.027037492
    614 : 0.04667454
    615 : 0.03902883
    616 : 0.031645115
    617 : 0.024575278
    618 : 0.041405562
    619 : 0.043281656
    620 : 0.03259573
    621 : 0.045636415
    622 : 0.038374618
    623 : 0.037076816
    624 : 0.060351457
    625 : 0.04088253
    626 : 0.07273529
    627 : 0.060353585
    628 : 0.037837572
    629 : 0.06006176
    630 : 0.059782706
    631 : 0.039284036
    632 : 0.042368222
    633 : 0.061792936
    634 : 0.04852204
    635 : 0.03687985
    636 : 0.07454695
    637 : 0.040446464
    638 : 0.048590116
    639 : 0.06772984
    640 : 0.05141519
    641 : 0.039789584
    642 : 0.044170357
    643 : 0.06999623
    644 : 0.034520544
    645 : 0.045540407
    646 : 0.042383093
    647 : 0.03365759
    648 : 0.021405222
    649 : 0.035889767
    650 : 0.036195457
    651 : 0.027589567
    652 : 0.04239747
    653 : 0.039578654
    654 : 0.028770007
    655 : 0.030080456
    656 : 0.050031155
    657 : 0.029260868
    658 : 0.034224443
    659 : 0.020916695
    660 : 0.03074026
    661 : 0.035108253
    662 : 0.035204493
    663 : 0.023221016
    664 : 0.0318839
    665 : 0.0320149
    666 : 0.025719002
    667 : 0.025805574
    668 : 0.028186977
    669 : 0.034081545
    670 : 0.048085738
    671 : 0.031097393
    672 : 0.03472594
    673 : 0.024711333
    674 : 0.032001123
    675 : 0.021611901
    676 : 0.037393764
    677 : 0.019381454
    678 : 0.026315255
    679 : 0.02667943
    680 : 0.021096552
    681 : 0.015127234
    682 : 0.03914332
    683 : 0.031361226
    684 : 0.021078518
    685 : 0.045183096
    686 : 0.027926063
    687 : 0.028400354
    688 : 0.021537596
    689 : 0.026974322
    690 : 0.032721475
    691 : 0.025840454
    692 : 0.022723919
    693 : 0.030373514
    694 : 0.027074365
    695 : 0.038588595
    696 : 0.019219844
    697 : 0.018751921
    698 : 0.019575367
    699 : 0.02089243
    700 : 0.012451103

    Saving weights

    Epoch # :  8
    701 : 0.021687936
    702 : 0.018156799
    703 : 0.019069392
    704 : 0.024869008
    705 : 0.021282775
    706 : 0.013152667
    707 : 0.012999469
    708 : 0.021141393
    709 : 0.022853188
    710 : 0.022628894
    711 : 0.013591609
    712 : 0.019952398
    713 : 0.023527546
    714 : 0.020037923
    715 : 0.023649784
    716 : 0.01514801
    717 : 0.022660999
    718 : 0.024182048
    719 : 0.026833503
    720 : 0.017029552
    721 : 0.01890941
    722 : 0.018819181
    723 : 0.023298416
    724 : 0.017434366
    725 : 0.019605149
    726 : 0.022318969
    727 : 0.021990411
    728 : 0.026353426
    729 : 0.023402438
    730 : 0.02011758
    731 : 0.018073458
    732 : 0.01866623
    733 : 0.036128055
    734 : 0.018028194
    735 : 0.031982526
    736 : 0.020512795
    737 : 0.02091242
    738 : 0.01729165
    739 : 0.019218981
    740 : 0.025949769
    741 : 0.013242643
    742 : 0.018119847
    743 : 0.018498959
    744 : 0.017910143
    745 : 0.015983991
    746 : 0.02003546
    747 : 0.026783576
    748 : 0.023054764
    749 : 0.02125796
    750 : 0.026635446
    751 : 0.016578453
    752 : 0.0120914765
    753 : 0.028330885
    754 : 0.014503027
    755 : 0.016883217
    756 : 0.0147133
    757 : 0.016482448
    758 : 0.01295083
    759 : 0.011887768
    760 : 0.017412286
    761 : 0.017432526
    762 : 0.014522076
    763 : 0.018125582
    764 : 0.01527135
    765 : 0.022896195
    766 : 0.033599176
    767 : 0.015266767
    768 : 0.024155777
    769 : 0.03737206
    770 : 0.025127921
    771 : 0.041732177
    772 : 0.018059555
    773 : 0.024936654
    774 : 0.024879891
    775 : 0.013694754
    776 : 0.015823958
    777 : 0.03468734
    778 : 0.015017441
    779 : 0.026725011
    780 : 0.027229277
    781 : 0.025583792
    782 : 0.029041352
    783 : 0.027529519
    784 : 0.018314665
    785 : 0.028311983
    786 : 0.023246178
    787 : 0.02675581
    788 : 0.053435303
    789 : 0.013630798
    790 : 0.026709728
    791 : 0.027024984
    792 : 0.016741473
    793 : 0.016449234
    794 : 0.022358602
    795 : 0.018893803
    796 : 0.011198228
    797 : 0.020461403
    798 : 0.024748158
    799 : 0.016092092
    800 : 0.032116514

    Saving weights

    Epoch # :  9
    801 : 0.023527348
    802 : 0.010949402
    803 : 0.019831575
    804 : 0.021879008
    805 : 0.018150676
    806 : 0.018971229
    807 : 0.024053812
    808 : 0.021102972
    809 : 0.015863054
    810 : 0.023519848
    811 : 0.023070032
    812 : 0.026120929
    813 : 0.02378812
    814 : 0.031062154
    815 : 0.017209096
    816 : 0.023328515
    817 : 0.03255829
    818 : 0.023766179
    819 : 0.015625024
    820 : 0.020562897
    821 : 0.049517382
    822 : 0.015665242
    823 : 0.03611715
    824 : 0.036744185
    825 : 0.015715063
    826 : 0.019078137
    827 : 0.029564833
    828 : 0.03045818
    829 : 0.021136235
    830 : 0.021629859
    831 : 0.020041553
    832 : 0.023357593
    833 : 0.016136
    834 : 0.013937189
    835 : 0.01734501
    836 : 0.019771762
    837 : 0.016841218
    838 : 0.01536422
    839 : 0.017983807
    840 : 0.018204879
    841 : 0.02033709
    842 : 0.023387741
    843 : 0.012951067
    844 : 0.015609992
    845 : 0.015595074
    846 : 0.020676674
    847 : 0.019395785
    848 : 0.017722672
    849 : 0.016651006
    850 : 0.025165152
    851 : 0.015526069
    852 : 0.02082716
    853 : 0.018446414
    854 : 0.013976221
    855 : 0.016713766
    856 : 0.015360777
    857 : 0.01498841
    858 : 0.02231152
    859 : 0.015597709
    860 : 0.015972357
    861 : 0.020932779
    862 : 0.014828679
    863 : 0.010003715
    864 : 0.023374233
    865 : 0.018153435
    866 : 0.01432657
    867 : 0.017923357
    868 : 0.01503726
    869 : 0.01836666
    870 : 0.02018271
    871 : 0.023310866
    872 : 0.01864164
    873 : 0.018086938
    874 : 0.013175893
    875 : 0.01396496
    876 : 0.022119215
    877 : 0.01871619
    878 : 0.017512083
    879 : 0.017614363
    880 : 0.018195935
    881 : 0.012661024
    882 : 0.01820031
    883 : 0.02199453
    884 : 0.010833663
    885 : 0.020424845
    886 : 0.0184608
    887 : 0.009907838
    888 : 0.021891676
    889 : 0.022460546
    890 : 0.022284722
    891 : 0.014578708
    892 : 0.017027227
    893 : 0.024437241
    894 : 0.014151578
    895 : 0.01371657
    896 : 0.011172482
    897 : 0.019281097
    898 : 0.016283643
    899 : 0.01758026
    900 : 0.017860396

    Saving weights

    Epoch # :  10
    901 : 0.023983361
    902 : 0.014110123
    903 : 0.017544664
    904 : 0.018290961
    905 : 0.013644686
    906 : 0.023201706
    907 : 0.012910295
    908 : 0.013804574
    909 : 0.010367661
    910 : 0.016588721
    911 : 0.013219259
    912 : 0.009908446
    913 : 0.022870334
    914 : 0.012612653
    915 : 0.011707573
    916 : 0.009927098
    917 : 0.010253669
    918 : 0.013215917
    919 : 0.01884599
    920 : 0.013364163
    921 : 0.014802155
    922 : 0.019709343
    923 : 0.015578944
    924 : 0.013616281
    925 : 0.015521779
    926 : 0.012936633
    927 : 0.012167938
    928 : 0.014166577
    929 : 0.013011998
    930 : 0.016763898
    931 : 0.018545182
    932 : 0.01663102
    933 : 0.016273562
    934 : 0.010520632
    935 : 0.0105387345
    936 : 0.009943314
    937 : 0.011543599
    938 : 0.010952039
    939 : 0.011538068
    940 : 0.012453527
    941 : 0.011754112
    942 : 0.016636811
    943 : 0.017222516
    944 : 0.009324293
    945 : 0.009084051
    946 : 0.0143687315
    947 : 0.010870088
    948 : 0.012142699
    949 : 0.016324442
    950 : 0.009925196
    951 : 0.0100024
    952 : 0.0126454355
    953 : 0.0146793565
    954 : 0.013418141
    955 : 0.013895083
    956 : 0.01397774
    957 : 0.013729873
    958 : 0.015033621
    959 : 0.015328206
    960 : 0.019933945
    961 : 0.014806671
    962 : 0.018491514
    963 : 0.017737335
    964 : 0.013863844
    965 : 0.013754041
    966 : 0.018179528
    967 : 0.013488879
    968 : 0.013930746
    969 : 0.008715766
    970 : 0.007628554
    971 : 0.016496891
    972 : 0.012624218
    973 : 0.0097497795
    974 : 0.014186646
    975 : 0.01849011
    976 : 0.013620214
    977 : 0.016832981
    978 : 0.0137945535
    979 : 0.016127512
    980 : 0.0133153675
    981 : 0.009538996
    982 : 0.014279413
    983 : 0.012188692
    984 : 0.008317082
    985 : 0.014955339
    986 : 0.015025271
    987 : 0.016098987
    988 : 0.011772374
    989 : 0.017549185
    990 : 0.013207734
    991 : 0.013621126
    992 : 0.011447512
    993 : 0.012098037
    994 : 0.014460028
    995 : 0.01570248
    996 : 0.012644585
    997 : 0.026523234
    998 : 0.009327965
    999 : 0.021500934
    1000 : 0.014526849

    Saving weights

    Wall time: 3min 10s

```python
# remove noise from initial epochs
loss_history_plot = list(filter(lambda x: x < 0.25, loss_history))

plt.figure(figsize=(12, 5))
plt.plot(loss_history_plot)
```

# Restoring custom models

Because of the added complexity of saving hybrid models, we need to have a slightly different method to load the weights of the custom model from the checkpoint.

```python
if os.path.exists(CHECKPOINTS_DIR):
    ckpt_path = tf.train.latest_checkpoint(CHECKPOINTS_DIR)
    print("Checkpoint path : ", ckpt_path)

    with tf.device(_device):
        model = EagerJANETModel(input_dim, NUM_UNITS, num_outputs=1, num_timesteps=TIME_STEPS, output_activation='linear')

        # model.compile(tf.train.AdamOptimizer(), loss='mse')

        # this is where you need to call the model at least once,
        # so that all of its variables can be properly restored
        zeros = tf.zeros((1, TIME_STEPS, 2))
        model(zeros)

        # restore the weights
        # see the difference as to how we pass the additional weight matrices to be loaded
        tfe.Checkpoint(model=model).restore(ckpt_path)

        # predict one batch to ensure the weights are correctly loaded
        generator = batch_generator()

        losses = []
        for i in range(20):
            inputs, outputs = next(generator)

            preds = model(inputs)
            loss_val = loss(outputs, preds)

            losses.append(loss_val.numpy())

    print("Final average predicted error (should be less than 0.03) : ", np.mean(losses))
```

    Checkpoint path :  checkpoints_addition/addition_janet.ckpt-1
    Final average predicted error (should be less than 0.03) :  0.012852971

```python

```


# Allocation Agent: Monte Carlo Tree Search (MCTS)

A full review of MCTS is available in [Browne et al., 2012](http://ccg.doc.gold.ac.uk/ccg_old/papers/browne_tciaig12_1.pdf), and has been the basis of the below work.

Broadly speaking MCTS randomly explores a decision space, building a tree of connected moves and their associated rewards. Over multiple iterations the MCTS determines which decisions are most likely to result in a positive outcome. Ultimately the best decision is the one that provides the best long term reward, the definition of which depends on the specific domain. For example, when creating a MCTS agent to play a board game, the long term reward would be a 0 or 1 depending on whether you won or lost the game after making the current move.

In the context of bed allocation we do not have a natural end state. Therefore the long term reward is determined as the total reward incurred after N time has passed according to the equation below. Here $R_{n}$ represents a reward associated with  the state of the hospital at a given time step, $\gamma \epsilon [0, 1]$ is the discount factor. The reward associated with a hospital state is $1 - total penalties$ incurred. The first term in the equation, $R_{1}$, is the immediate reward associated with the current allocation i.e., the greedy allocation score, and the subsequent terms are rewards associated with future states, where the discount factor weighs the relative importance of these future states against the current state.
$$
\begin{align}
R = R_{1} + \gamma R_{2} +\gamma^{2}R_{3} ... +\gamma^{N}R_{N}
\end{align}
$$

In this notebook we demonstrate how to run the MCTS allocation algorithm in a simplified virtual hospital; see `src/agent/mcts` and `src/agent/simulation` for the relevant code. More details of the implementation are provided at the end of this notebook.

## 1. Import required modules

_Note:_ you will need to first install the module as per the instructions in the main README, and run a notebook server from within the same virtual environment to have access to the `hospital` submodules.

```python
import cloudpickle
import copy
import time
import random

import warnings
warnings.filterwarnings('ignore')

from tqdm.notebook import tqdm

import agent.utils as utils
import agent.policy as policy
import agent.run_mcts as mcts
from hospital.people import Patient
from hospital.building import Hospital, MedicalWard, SurgicalWard, Room
from hospital.equipment.bed import Bed
import hospital.restrictions.ward as R
```

## 2. Hospital Environment

The MCTS implementation can take a long time to run, and have high memory requirements. To demonstrate how it can be run we create a simplified scenario with a hospital containing 2 wards (1 medical, 1 surgical), and 5 beds each. The Medical ward will have a restriction for not allowing surgical patients and vice versa.

```python
beds = [Bed(name=f"B00{i}") for i in range(10)]
wards = [
    MedicalWard(name="MedicalWard", rooms=[Room(name="R000", beds=beds[:5])]),
    SurgicalWard(name="SurgicalWard", rooms=[Room(name="R001", beds=beds[5:])]),
]
h = Hospital(name="Hospital", wards=wards)

# Add ward restrictions
h.wards[0].restrictions = [R.NoSurgical(10)]
h.wards[1].restrictions = [R.NoMedical(5)]

# Populate at 50%
policy.populate_hospital(h, occupancy=0.5)
h.render()
```

```python
# you will also need to normalise the penalties to lie between 0 and 1
utils.normalise_ward_penalties(h)
```

## 3. Create patients to admit to the hospital

To generate realistic patients utilise the PatientSampler class which takes a day of week and an hour of the day and returns a forecast for the number of patients estimated to arrive each hour, `N`. By default, the patients are synthesied with random data where the distribution for each attributes is informed by aggregated historic data. If historic patient data is available then a pool of historic patients can be saved and setting `historic=True` will return more accurate marginal distributions accross all patient attributes. Otherwise, the returned number of patients `N` is a random number, and the returned patient attributes are randomly generated. See `src/forecasting/patient_sampler`.

```
sampler = PatientSampler("monday", 9)
forecast_window=2
forecasted_patients = sampler.sample_patients(forecast_window=forecast_window, num_samples=1)

# we can unpack the above structure into a list of lists
# each sublist represents an hour of patients
arrivals = []
for _, patients in forecasted_patients[0].items():
     arrivals.append(patients)
```

We will instead use the code below to create a simplified list of arriving patients, by initialising patients with the default value for most fields, such that there won't be any additional patient level restrictions (e.g., patient needs side room for immunosuppression). The arrivals list will still have the same structure as above.

```python
def generate_random_name() -> str:
    """
    2 random letters + 2 random numbers. As a string.
    """
    letters = "abcdefghijklmnopqrstuvwxyz"
    digits = "0123456789"
    characters = random.choices(letters, k=2) + random.choices(digits, k=2)
    name = "".join(characters)
    return name

def generate_simple_patient() -> Patient:
    """Returns a random patient, without any patient level restrictions."""

    # department and specialty
    department = ["medicine", "surgery"][utils.bernoulli(0.5)]

    # Patient, all other attributes are default=False
    patient = Patient(
        name=generate_random_name(),
        sex=["male", "female"][utils.bernoulli(0.5)],
        department=department
    )

    return patient

def generate_arrivals(max_per_hour: int, forecast_window: int) -> list:
    """Creates a list of arriving patients for each hour of in the forecast window"""
    arrivals =[]
    for hours in range(forecast_window):
        arrivals.append([generate_simple_patient() for _ in range(random.randint(0, max_per_hour))])

    return arrivals

```

```python
# patient currently being allocated
patient = Patient(
    name="patient_0",
    sex="female",
    department="medicine",
)

# 'forecasted' arrivals
arrivals = generate_arrivals(3, 4)

print(f"Incoming patients per hour: {[len(l) for l in arrivals]}")

# now we insert the patient we are currently trying to
# allocate as the first patient, as time t=1.
arrivals = [[patient]] + arrivals
```

Now we can run the MCTS as below. The output in `mcts_output` is a list of dictionaries where each dictionary represents a possible allocation for patient_0, and the results of the tree search for that option, further details of the tree search are provided below. The items of each dictionary are:
- `action`, the allocation of the current patient into a bed. It is possible to consider allocating multiple patients in a timestep, as such, the action is represented by a dictionary of bed_name:patient pairs, for each patient at t=0 in the tree search. In our example there is just one patient.
- `score`, the immediate pentaly associated with allocating the current patient(s) to the suggested bed(s).
- `violated_restrictions`, the set of restrictions (if any), violated by assigning the patient(s) to the suggested bed(s).
- `ucb_score`, the tree policy score associated with the suggested allocation. See below for more details on UCB score.
- `visit_count`, the number of times the node representing the suggested allocation was visited during tree search.

There are several potential strategies for determining what the best allocation option is. The one with the lowest `score` is equivalent to making a greedy optimisation that selects the best bed given the current circumstances of the hospital; whereas, making a choice between the highest `ucb_score`, or `visit_count` or a weighted average of both provides the best allocation according to the MCTS.

```python
h_copy = copy.deepcopy(h)
t = time.time()
mcts_node = mcts.run_mcts(
        h_copy,
        arrivals,
        discount_factor=0.9,
        n_iterations=100,
)
elapsed = time.time() - t
mcts_output = mcts.construct_mcts_output(h_copy, mcts_node, patient)
```

```python
print(f"Time taken to compute suggestions: {round(elapsed,2)}s")
mcts_output
```

# 4. Details on implementation

Below we describe the four stages of the MCTS algorithm as they are specifically implemented for the bed allocation agent. The implementation utilises the `anytree` library to build a tree structure, where each node represents a specific state of the hospital and each level of the tree represents a time step. Time steps are incremented in hours and connected to the number of forecasted admissions arriving each hour. The input to the treesearch is a queue of patients arriving at each time step, with the current patient to be allocated (`t=0`) as the first entry in this queue, and the current state of the hospital as the root node to search from.

Selection

Starting at the root node a child node is selected. The root node represents the current state of the hospital at time $t=0$ where the state is defined by the patients that are currently occupying beds and the set of empty beds.

In the first iteration of the tree search, the algorithm selects the root node and moves to the expansion step. In subsequent iterations it will traverse from the root and choose one of the child nodes according to the tree policy. The tree policy is the UCB score, in which $\overline{R}$ is the mean reward from visiting that node, $N_{pi}$ is the number of times its parent node was visited and $N_{i}$ is the number of times the node itself has been visited. The first term encourages the algorithm to select nodes that have previously resulted in good outcomes, while the second encourages the algorithm to explore options that it hasn’t visited as often.
$$
\begin{align}
UCB = \overline{R} + \sqrt{\frac{2 \log{N_{p}}}{N_{i}}}
\end{align}
$$

Expansion

Once a node is selected, a child node is attached to represent each possible decision state for that time step.

For example, we have a hospital with 4 beds, 2 are occupied and 2 are available. We are currently trying to allocate a patient P1. As there are two possible decisions, two nodes can be attached that represent 1) allocating P1 to the first available bed and 2) allocation P1 to the second available bed.

As we progress through the tree search, we may encounter time steps where multiple patients have arrived. In such cases, a node is expanded for each possible combination of patients to available beds.

Simulation

From one of the attached children we then simulate a future. The simulation stage involves the following steps:

Each patient currently within the hospital has a length of stay attribute, and an expected length of stay attribute. At the start of the simulation step, the length of stay counters are incremented by one.

Then a discharge model is applied to discharge existing patients. The probability of being discharged increases according to proximity to your expected length of stay.

The patients arriving in the given time-step are then assigned to beds according to the default policy. The default policy is a random assignment of patients to beds to available beds.

Backpropagation

The total penalty of the hospital is calculated after the simulation step. This is the sum of all penalties for each broken restriction within the hospital. We then backpropagate this score up the tree to distribute the outcome across all decisions along the currently explored path.

This stage updates the UCB score (tree policy score) and visit count for each node that was traversed along the current decision pathway. If the result of the simulation was good, the UCB scores of each node will have increased, making it more likely that future iterations of the tree search will select these nodes again. The reverse is also true. In this manner MCTS is more tractable than a completely random search of the possible decision pathways as it more frequently visits the most promising options during the selection stage.

The above procedure is repeated multiple times until a maximum number of iterations have been reached. At this point the tree object is returned and the best child node of the root is selected as the optimal allocation for the current patient. There are several potential strategies for determining what the best node is. In the current implementation we choose the node that has the highest visit count. Alternative approaches such as choosing the node with the highest UCB score or some balance of the two, and how this affects outcomes, remain to be explored in future work.

### Limitations
In the above implementation we take a single sample from the demand forecast and use this as a fixed version of the future. This means that the future within the treesearch is deterministic and significantly reduces the search space, and branching of the tree, allowing the algorithm to find a recommendation in a more tractable timeframe. However, a single sample from the forecast represents one of the possible futures. To truly capture the variability of incoming patients, we envisage a strategy where multiple simultaneous tree searches are implemented, each using a separate sample from the forecasted admissions. These could be run in parallel to increase runtime efficiency, and the final suggested allocation would be the bed that has the average highest ranking across the ensemble of tree searches. The efficacy of this strategy and alternative approaches to dealing with non deterministic search spaces remain to be explored.

Despite fixing the set of arrivals within a tree-search we can still experience an intractable amount of branching that makes the current implementation of MCTS unsuitable for operational use. For example, if there are just 4 empty beds in the hospital, and 9 arriving patients within a time step, the tree expands into 840 possible permutations of patients to beds. With multiple time steps into the future this can compound and result in either memory issues or extremely long compute times. Further work is needed to explore engineering strategies that can make MCTS more operationally feasible.


# 1. Teoria de Grafos

## 1.1. Definições Básicas

A Teoria de Grafos é um ramo da matemática que estuda as relações entre os objetos de um determinado conjunto. Um **grafo** é uma estrutura composta por **vértices** (nós) e **arestas** (arcos), de tal maneira que diversos problemas podem ser formulados utilizando esta estrutura.

Formalmente falando, um grafo $G=(V,E)$ é definido da seguinte forma:
* $V$ é o conjunto finito de elementos denominados vértices e
* $E$ é o conjunto finito de pares **não ordenados** de vértices chamados de arestas. Na prática, uma aresta conecta dois vértices.

A Figura 1 apresenta um exemplo de grafo, em que $V=\{A, B, C, D\}$ e $E=\{(A,B),(A,C),(B,D)\}$. Como mencionado anteriormente, $E$ é um conjunto não ordenado de vértices, o que significa que as arestas $e_1=(A,B)$ e $e_2=(B,A)$ representam a mesma relação e, por isso, apenas uma delas é incluída no conjunto $E$.

Figura 1: Exemplo de um grafo.

Algumas definições adicionais incluem: seja $e=(A,B)$ uma aresta tal que $e\in E$. Dizemos que $A$ e $B$ são os **extremos** desta aresta, e que $A$ e $B$ são vértices **adjacentes**. Ademais, temos que o **tamanho** do grafo $G$ é dado por $\left|V\right|+\left|E\right|$, em que $\left|V\right|$ e $\left|E\right|$ denotam, respectivamente, o número de vértices e arestas. No exemplo da Figura 1, temos que $\left|V\right|=4$ e $\left|E\right|=3$.

Dizemos que um grafo é **simples** (Figura 1, por exemplo) quando o mesmo não possui **laços** ou **arestas múltiplas**. Um laço é uma aresta com extremos idênticos, ao passo que arestas múltiplas são definidas como sendo duas ou mais arestas com o mesmo par de extremos, conforme demonstrado na Figura 2.

Figura 2: Exemplo de laços e arestas múltiplas em um grafo.

Dizemos que $H=(V^\prime,E^\prime)$ é um **subgrafo** de $G=(V,E)$ quando $V^\prime \subseteq V$ e $E^\prime \subseteq E$. Já o grafo $J=(V,E^\prime)$ é dito ser subgrafo **gerador** de $G$ quando $E^\prime \subseteq E$. A Figura 3 ilustra esses conceitos.

<image>

(a)

(b)

(c)

Figura 3: (a) grafo $G$, (b) subgrafo $J$ e (c) subgrafo gerador $H$.

</image>

Já com relação ao **grau** de um vértice $v\in V$, denotado por $d(v)$, o mesmo diz respeito à quantidade de nós adjacentes à $v$, sendo que laços contam duas vezes. Na Figura 4, por exemplo, temos que $d(A) = 2$, $d(B) = 4$, $d(C) = 1$ e $d(D) = 1.$

Figura 4: Exemplo de um grafo para contagem de graus de seus nós.

De acordo com o teorema conhecido por **Handshake Lemma**, temos a seguinte afirmação:

\begin{equation}
\tag{1}
\sum_{v\in V}d(v)=2\left|E\right|.
\end{equation}
Vamos tentar provar por indução? A prova será realizada de acordo com o **número de vértices** no conjunto $V$. Conforme vimos em aulas anteriores, uma prova por indução possui os seguintes passos:

1. Base: seja $G=(V,E)$ um grafo com apenas um vértice $v$, ou seja, $\left|V\right|=1$. Neste caso, existem duas possibilidades: (i) $d(v) = 0$ e (ii) $d(v) = 2$. Na primeira opção, temos que o número de arestas é igual à zero, ou seja, $\left|E\right|=0$. Assim sendo, a Equação 1 é verdadeira, ou seja, $\sum_{v\in V}d(v)=2\left|E\right|=2.0=0$. Para o segundo caso, temos que $G$ possui um laço no vértice $v$, ou seja, $d(v)=2$. Neste caso, a Equação 1 também é verdadeira, ou seja, $\sum_{v\in V}d(v)=2\left|E\right|=2.1=2$.

2. Hipótese de Indução: seja $G=(V,E)$ um grafo tal que a Equação 1 seja verdadeira.

```python

```