Rat model silently returns invalid predictions on TensorFlow >= 2.16 (Keras 3)

[NicoKiaru]

(Disclaimer: AI generated)

Summary

The rat model returns near-meaningless predictions when DeepSlice runs on TensorFlow >= 2.16, i.e. whenever tf.keras resolves to Keras 3. It fails silently: no exception, no warning, just bad coordinates. The mouse model is completely unaffected (bit-identical output under both Keras versions).

Since setup.py/pyproject pins tensorflow unpinned, a fresh pip install DeepSlice today installs TF 2.21, so I think anyone who has tried the rat model recently has been hitting this. It may account for part of the "the rat model is weak" reports in #68 and #78.

Why only the rat model

initialise_network() builds the two species differently:

if species == "rat":
    inputs = Input(shape=(299, 299, 3))
    base_model_layer = base_model(inputs, training=True)   # <-- nested call with training=True
    ...
    model = Model(inputs=inputs, outputs=output_layer)
else:
    model = Sequential()
    model.add(base_model)
    ...

Keras 3 does not treat that nested training=True the way Keras 2 did. A visible symptom: under Keras 3 the first DSModel("rat") created in a process and the second one give different predictions for the same images, and the second matches what you get by building the model with training=False. Neither is correct. Under Keras 2 the result is stable and correct.

Reproducer

49 coronal sections, same JPEGs in all runs, predict(folder, ensemble=False, section_numbers=True) then propagate_angles(). The sections are numbered in cutting order, so the predicted depth oy should increase monotonically along the series.

	Spearman(oy, section order)	median |u|	cv(|u|)	median |v|	spread of plane normal
TF 2.21, Keras 3	-0.06	188	0.48	305	0.46
TF 2.21 + tf-keras, TF_USE_LEGACY_KERAS=1	+0.97	618	0.03	590	0.000

Under Keras 3 the anchoring vectors are ~4x too short, vary by 50% from slice to slice, the plane normal flips sign between slices, and depth ordering is random. Under Keras 2 everything is consistent and the depth ordering is recovered.

Workaround

Installing the Keras 2 backport and enabling it before TensorFlow is imported fixes it completely:

pip install tf-keras==2.21.0

import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"   # must precede any tensorflow import
from DeepSlice import DSModel

Suggestion

Either pin tensorflow<2.16 in the requirements, or declare tf-keras as a dependency and set TF_USE_LEGACY_KERAS inside DeepSlice/__init__.py before TensorFlow is imported, so users get correct results out of the box. A warning when keras.__version__ starts with 3 would also have saved a lot of head-scratching.

I am happy to open a PR if that would help.

Found while debugging the DeepSlice integration in ABBA; we have applied the tf-keras workaround on our side.

[PolarBean]

Oh great, thank you for finding this. A PR would be very welcome!