Tensorflow_hub: ray.exceptions.RaySystemError: System error: Unknown layer: 'KerasLayer'

Ray Libraries (Data, Train, Tune, Serve)
1

Using Tensorflow Hub, to use Google’s NNLM: following this example: TensorFlow Hub

Getting the following error:

Failure # 1 (occurred at 2023-07-05_15-44-12)
e[36mray::_Inner.train()e[39m (pid=40680, ip=127.0.0.1, actor_id=66171e7658f7e78b22035c4001000000, repr=TensorflowTrainer)
File “python\ray_raylet.pyx”, line 1434, in ray._raylet.execute_task
File “python\ray_raylet.pyx”, line 1438, in ray._raylet.execute_task
File “python\ray_raylet.pyx”, line 1378, in ray._raylet.execute_task.function_executor
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray_private\function_manager.py”, line 724, in actor_method_executor
return method(__ray_actor, *args, **kwargs)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray\util\tracing\tracing_helper.py”, line 464, in _resume_span
return method(self, *_args, **_kwargs)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray\tune\trainable\trainable.py”, line 389, in train
raise skipped from exception_cause(skipped)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray\train_internal\utils.py”, line 54, in check_for_failure
ray.get(object_ref)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray_private\auto_init_hook.py”, line 18, in auto_init_wrapper
return fn(*args, **kwargs)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray_private\client_mode_hook.py”, line 103, in wrapper
return func(*args, **kwargs)
File “C:\ProgramData\Anaconda3\envs\forecast_bmsc\lib\site-packages\ray_private\worker.py”, line 2542, in get
raise value
ray.exceptions.RaySystemError: System error: Unknown layer: ‘KerasLayer’. Please ensure you are using a keras.utils.custom_object_scope and that this object is included in the scope. …
traceback: Traceback (most recent call last):

Solution is to specify the custom object, as discussed in this thread: python - Unknown layer: KerasLayer when i try to load_model - Stack Overflow

my_reloaded_model = tf.keras.models.load_model(
       (path),
       custom_objects={'KerasLayer':hub.KerasLayer}
)

But… I’ve read through all the Ray doc, there’s no way to apply this fix within Ray? Help?

2

Worked out the solution to this problem. Posting here to help others.

add the following declaration atop:

@keras.utils.register_keras_serializable(package=“HubLayers”)
class KerasLayer(hub.KerasLayer):
…

Then use defined KerasLayer instead of hub:

embed_layer = KerasLayer(
embedding_url,
input_shape=(),
output_shape=(50,),
dtype=tf.string,
trainable=True,
)