use tf.shape for `batch_size` for h_/c_init

[jsadler2]

@jzwart - this allows the models to run when compiled without run_eagerly=True

closes #108

[jzwart]

Oh good to know. So produces the same when run_eagerly=True or not ? If I remember, run_eagerly=True can slow the model down, right? So if this works, then we shouldn't have to set run_eagerly=True ever

[SimonTopp]

Nice. Thanks @jsadler2 !

[jsadler2]

So produces the same when run_eagerly=True or not ? Let me test this

[jsadler2]

I just tested it and here are the outputs of the two versions.

Summary: when run_eagerly=True the model

• compiles a bit faster the in pretraining
• compiles a lot faster in finetuning
• runs ~3x slower per training step (year in our case)

So unless you are running only a few epochs, it's probably better to do run_eagerly=False (which is the default).

Detailed output

run_eagerly=True:

Epoch 1/10
21/21 [==============================] - 68s 223ms/step - loss: 282.3889
Epoch 2/10
21/21 [==============================] - 4s 201ms/step - loss: 235.5191
Epoch 3/10
21/21 [==============================] - 5s 219ms/step - loss: 210.9826
Epoch 4/10
21/21 [==============================] - 4s 211ms/step - loss: 177.4395
Epoch 5/10
21/21 [==============================] - 4s 202ms/step - loss: 153.2811
Epoch 6/10
21/21 [==============================] - 5s 224ms/step - loss: 149.5416
Epoch 7/10
21/21 [==============================] - 4s 201ms/step - loss: 142.6568
Epoch 8/10
21/21 [==============================] - 4s 215ms/step - loss: 131.4719
Epoch 9/10
21/21 [==============================] - 4s 191ms/step - loss: 138.4801
Epoch 10/10
21/21 [==============================] - 4s 190ms/step - loss: 141.5008
Epoch 1/10
21/21 [==============================] - 5s 222ms/step - loss: 155.2064
Epoch 2/10
21/21 [==============================] - 4s 204ms/step - loss: 114.1123
Epoch 3/10
21/21 [==============================] - 5s 222ms/step - loss: 115.4744
Epoch 4/10
21/21 [==============================] - 4s 204ms/step - loss: 103.2803
Epoch 5/10
21/21 [==============================] - 4s 211ms/step - loss: 101.4940
Epoch 6/10
21/21 [==============================] - 5s 222ms/step - loss: 100.5679
Epoch 7/10
21/21 [==============================] - 4s 206ms/step - loss: 108.1179
Epoch 8/10
21/21 [==============================] - 4s 206ms/step - loss: 104.9283
Epoch 9/10
21/21 [==============================] - 5s 225ms/step - loss: 98.3793
Epoch 10/10
21/21 [==============================] - 4s 208ms/step - loss: 96.0975

run_eagerly=False

Epoch 1/10
21/21 [==============================] - 73s 78ms/step - loss: 271.2268
Epoch 2/10
21/21 [==============================] - 2s 81ms/step - loss: 223.4635
Epoch 3/10
21/21 [==============================] - 2s 84ms/step - loss: 185.6641
Epoch 4/10
21/21 [==============================] - 2s 84ms/step - loss: 154.2734
Epoch 5/10
21/21 [==============================] - 2s 83ms/step - loss: 148.8328
Epoch 6/10
21/21 [==============================] - 2s 72ms/step - loss: 143.4130
Epoch 7/10
21/21 [==============================] - 1s 70ms/step - loss: 151.3037
Epoch 8/10
21/21 [==============================] - 1s 70ms/step - loss: 142.6057
Epoch 9/10
21/21 [==============================] - 2s 73ms/step - loss: 148.6824
Epoch 10/10
21/21 [==============================] - 2s 80ms/step - loss: 143.5335
Epoch 1/10
21/21 [==============================] - 35s 88ms/step - loss: 144.0985
Epoch 2/10
21/21 [==============================] - 2s 89ms/step - loss: 117.7280
Epoch 3/10
21/21 [==============================] - 2s 89ms/step - loss: 125.7732
Epoch 4/10
21/21 [==============================] - 2s 88ms/step - loss: 106.1030
Epoch 5/10
21/21 [==============================] - 2s 88ms/step - loss: 106.6340
Epoch 6/10
21/21 [==============================] - 2s 72ms/step - loss: 113.5239
Epoch 7/10
21/21 [==============================] - 1s 71ms/step - loss: 108.2322
Epoch 8/10
21/21 [==============================] - 1s 70ms/step - loss: 112.8934
Epoch 9/10
21/21 [==============================] - 1s 70ms/step - loss: 124.7646
Epoch 10/10
21/21 [==============================] - 1s 70ms/step - loss: 99.5339