Interrupting training crashes Python kernel when training on GPU for the first time with MirroredStrategy

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
• TensorFlow installed from (source or binary): binary (pip install tensorflow)
• TensorFlow version (use command below): 2.4.3
• Python version: 3.7
• GPU model and memory: Nvidia Tesla V100

Describe the problem.

When using multi-GPU training in a Jupyter Notebook:

strategy = tf.distribute.MirroredStrategy()
model_ = ...
with strategy.scope():
    model = model_.compile(loss='categorical_crossentropy', optimizer='adam', metrics=["accuracy"])

Then launching training:

model.fit(...)

Describe the current behavior.

• when initial training runs to the end, subsequent trainings can be interrupted fine
• when initial training itself is interrupted, the Jupyter kernel dies

Detailed step-by-step procedure:

1. compile the model with mirrored strategy
2. launch training with model.fit()
3. let the training go to its own end -> otherwise kernel dies if interrupted here
4. after initial training completed, re-compile the model with mirrored strategy
5. launch another training
6. interrupt this training -> works fine, kernel doesn’t die

If training was interrupted at step 3. then the kernel dies systematically.

Describe the expected behavior.

• when initial training is interrupted, the Jupyter kernel should not die

Source code / logs.

No error message — the Jupyter iPython kernel simply dies with no message.

Current workaround.

The current workaround I found is to run the first training with just 1 epoch, let it complete, then I can recompile the model and re-train at will with no crashes. Less than ideal, but it works.

[jvishnuvardhan]

@JivanRoquet Quick question to find whether the root-cause is keras or distribution strategy. I guess distribution strategy. Are you using Keras or TF model? Did you try without Mirrored strategy? Do you notice this issues without distribution strategy.

If you think this is more related to distribution strategy, please open this issue in TF repository https://github.com/tensorflow/tensorflow/issues as this keras repo is mainly for keras related issues. Thanks!

[JivanRoquet]

Sure thing. Closing here for now.