Layer or Model in invalid state TF 2.12 - bioimageio export - _compute_receptive_field

[esgomezm]

Describe the bug
I'm trying to export a StarDist model using the function export_bioimageio but I got an error "Your Layer or Model is in an invalid state". Unfortunately, this breaks the compatibility with the bioimage model zoo. Seems the error comes from the CSBDeep models as it fails when running model.keras_mode.predict rather than StarDist directly. This error was already reported here. Is there any workaround we could do?

To reproduce
I'm using the ZeroCostDL4Mic that installs TensorFlow 2.12 (the latest TF version given by google colab). Everything works well expect the export.

Data and screenshots

Loading network weights from 'weights_best.h5'.
Loading thresholds from 'thresholds.json'.
Using default values: prob_thresh=0.163235, nms_thresh=0.3.
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[/usr/local/lib/python3.9/dist-packages/stardist/models/base.py](https://localhost:8080/#) in _axes_tile_overlap(self, query_axes)
   1081         try:
-> 1082             self._tile_overlap
   1083         except AttributeError:

AttributeError: 'StarDist2D' object has no attribute '_tile_overlap'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
6 frames
[<ipython-input-14-5a9275ed7113>](https://localhost:8080/#) in <cell line: 33>()
     31 model = StarDist2D(None, name=QC_model_name, basedir=QC_model_path)
     32 modelfile = Path(os.path.join(QC_model_path, QC_model_name + 'bioimageio.zip'))
---> 33 export_bioimageio(model, modelfile, test_img)

[/usr/local/lib/python3.9/dist-packages/stardist/bioimageio_utils.py](https://localhost:8080/#) in export_bioimageio(model, outpath, test_input, test_input_axes, test_input_norm_axes, name, mode, min_percentile, max_percentile, overwrite_spec_kwargs)
    397         tmp_dir = Path(_tmp_dir)
    398         kwargs = _get_stardist_metadata(tmp_dir, model)
--> 399         model_kwargs = _get_weights_and_model_metadata(tmp_dir, model, test_input, test_input_axes, test_input_norm_axes, mode,
    400                                                        min_percentile=min_percentile, max_percentile=max_percentile)
    401         kwargs.update(model_kwargs)

[/usr/local/lib/python3.9/dist-packages/stardist/bioimageio_utils.py](https://localhost:8080/#) in _get_weights_and_model_metadata(outdir, model, test_input, test_input_axes, test_input_norm_axes, mode, min_percentile, max_percentile)
    243 
    244     assert all(s in (0, 1) for s in output_scale), "halo computation assumption violated"
--> 245     halo = model._axes_tile_overlap(output_axes.replace('b', 's'))
    246     halo = [int(np.ceil(v/8)*8) for v in halo]  # optional: round up to be divisible by 8
    247 

[/usr/local/lib/python3.9/dist-packages/stardist/models/base.py](https://localhost:8080/#) in _axes_tile_overlap(self, query_axes)
   1082             self._tile_overlap
   1083         except AttributeError:
-> 1084             self._tile_overlap = self._compute_receptive_field()
   1085         overlap = dict(zip(
   1086             self.config.axes.replace('C',''),

[/usr/local/lib/python3.9/dist-packages/stardist/models/base.py](https://localhost:8080/#) in _compute_receptive_field(self, img_size)
   1067         z = np.zeros_like(x)
   1068         x[(0,)+mid+(slice(None),)] = 1
-> 1069         y  = self.keras_model.predict(x)[0][0,...,0]
   1070         y0 = self.keras_model.predict(z)[0][0,...,0]
   1071         grid = tuple((np.array(x.shape[1:-1])/np.array(y.shape)).astype(int))

[/usr/local/lib/python3.9/dist-packages/keras/engine/training_v1.py](https://localhost:8080/#) in predict(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)
   1052                 that is not a multiple of the batch size.
   1053         """
-> 1054         self._assert_built_as_v1()
   1055         base_layer.keras_api_gauge.get_cell("predict").set(True)
   1056         self._check_call_args("predict")

[/usr/local/lib/python3.9/dist-packages/keras/engine/base_layer_v1.py](https://localhost:8080/#) in _assert_built_as_v1(self)
    904     def _assert_built_as_v1(self):
    905         if not hasattr(self, "_originally_built_as_v1"):
--> 906             raise ValueError(
    907                 "Your Layer or Model is in an invalid state. "
    908                 "This can happen for the following cases:\n "

ValueError: Your Layer or Model is in an invalid state. This can happen for the following cases:
 1. You might be interleaving estimator/non-estimator models or interleaving models/layers made in tf.compat.v1.Graph.as_default() with models/layers created outside of it. Converting a model to an estimator (via model_to_estimator) invalidates all models/layers made before the conversion (even if they were not the model converted to an estimator). Similarly, making a layer or a model inside a a tf.compat.v1.Graph invalidates all layers/models you previously made outside of the graph.
2. You might be using a custom keras layer implementation with custom __init__ which didn't call super().__init__.  Please check the implementation of <class 'keras.engine.functional.Functional'> and its bases.

Environment (please complete the following information):

• stardist==0.8.3
• csbdeep==0.7.3
• TensorFlow version 2.12.0
• Google Colab

[esgomezm]

@FynnBe @constantinpape @oeway @cfusterbarcelo

[uschmidt83]

Hi @esgomezm,

sorry for the late reply.

I already mentioned a workaround here, but I assume you want something that works with TensorFlow 2?

---

I think the TF model export has to be different in TensorFlow 2, but I don't even know what the format should look like. I need a specific test case, i.e. a consumer for the exported model that can tell me whether the export was correct.

If you already know the answers to these questions, you can write/prototype it yourself. It's a bit hacky, but I think you can override model.export_TF before calling export_bioimageio, something like this:

def my_new_export_TF(self, fname=None, single_output=True, upsample_grid=True):
    // new code that doesn't call csbdeep.utils.tf.export_SavedModel
    ...

model.export_TF = my_new_export_TF
export_bioimageio(model, ...)

[esgomezm]

Hi Uwe,

Thanks for the replay. The problem is not the export but running model.keras_model.predict for the execution of the halo in TensorFlow 2. I managed to get a workaround as follows:

def _my_compute_receptive_field(model, img_size=None):
    # TODO: good enough?
    from scipy.ndimage import zoom
    if img_size is None:
        img_size = tuple(g*(128 if model.config.n_dim==2 else 64) for g in model.config.grid)
    if np.isscalar(img_size):
        img_size = (img_size,) * model.config.n_dim
    img_size = tuple(img_size)
    # print(img_size)
    assert all(_is_power_of_2(s) for s in img_size)
    mid = tuple(s//2 for s in img_size)
    x = np.zeros((1,)+img_size+(model.config.n_channel_in,), dtype=np.float32)
    z = np.zeros_like(x)
    x[(0,)+mid+(slice(None),)] = 1
    #y  = self.keras_model.predict(x)[0][0,...,0]
    #y0 = self.keras_model.predict(z)[0][0,...,0]
    y  = model.predict(np.squeeze(x))[0]
    y0 = model.predict(np.squeeze(z))[0]

    grid = tuple((np.array(x.shape[1:-1])/np.array(y.shape)).astype(int))
    assert grid == model.config.grid
    y  = zoom(y, grid,order=0)
    y0 = zoom(y0,grid,order=0)
    ind = np.where(np.abs(y-y0)>0)
    return [(m-np.min(i), np.max(i)-m) for (m,i) in zip(mid,ind)]

model._tile_overlap = _my_compute_receptive_field(model)
export_bioimageio(model, modelfile, test_img)

[uschmidt83]

The problem is not the export but running model.keras_model.predict for the execution of the halo in TensorFlow 2.

I'm confused, this is the first time I'm hearing about this. So what's the actual problem again? I still don't understand.

[esgomezm]

Hi @uschmidt83

The problem is that with TensorFlow 2.12 model.keras_mode.predict does not run (please, see the error message I copied at the beginning of the issue). I'm not sure how model.predict runs (which is the one working) but something seems to be differente

[uschmidt83]

The problem is that with TensorFlow 2.12 model.keras_mode.predict does not run (please, see the error message I copied at the beginning of the issue). I'm not sure how model.predict runs (which is the one working) but something seems to be differente

In the end, model.predict is just calling model.keras_model.predict under the hood.

The problem is not the export but running model.keras_model.predict for the execution of the halo in TensorFlow 2.

I believe this is just the symptom, not the real cause of the model export problem. Hence, I'm confused why your fix via model._tile_overlap = _my_compute_receptive_field(model) even works.

I feel like we're not making much progress here. Can you make a minimal example for me to reproduce the problem? I.e. a few lines of code that trigger the problem and which version of TensorFlow to use.