Ошибка при использовании keras.saving.save_model: «UnliftableError: невозможно поднять тензор, поскольку он транзитивно ⇐ Python

[Anonymous]

Я работаю над собственной моделью многорукого бандита в tensorflow keras 2.14.1
Модель можно вызвать и она работает отлично, но всякий раз, когда я пытаюсь сохранить ее с помощью keras.saving.save_model

Код: Выделить всё

import numpy as np
import keras
import tensorflow as tf

class MyArm(tf.keras.layers.Layer):
def __init__(self, input_dim: int) -> None:
super().__init__()
_zero_init = tf.keras.initializers.Constant(0)
self.input_dim = input_dim
self.score_scalar = self.add_weight(name="score_scalar", dtype=tf.float32, initializer=_zero_init, trainable=False)

def call(self, inputs: tf.Tensor) -> tf.Tensor:
return self.score_scalar

class MyModel(tf.keras.Model):
def __init__(self, input_dim: int, n_arms: int, name: str = "Dummy") -> None:
super().__init__(name=name)
self.input_dim = input_dim
self.n_arms = n_arms
self.model_layers = [MyArm(self.input_dim) for arm in range(n_arms)]
self.current_arm = 0  # Defaults points to first arm
self._set_current_arm_score(1.0)  # Raise arm

def _set_current_arm_score(self, score: float):
self.model_layers[self.current_arm].score_scalar.assign(score)

def _cycle(self):
self._set_current_arm_score(0.0)  # Lower arm
self.current_arm = (self.current_arm + 1) % self.n_arms  # Points to next arm in the cycle
self._set_current_arm_score(1.0)  # Raise arm

def compute_score(self, inputs: np.ndarray) -> tf.Tensor:
return tf.Variable([self.score_arm(arm, inputs) for arm in range(self.n_arms)])

def score_arm(self, arm: int, inputs: np.ndarray) -> tf.Tensor:
return self.model_layers[arm](inputs)

def call(self, inputs: np.array) -> tf.Tensor:
action = tf.argmax(self.compute_score(inputs))
self._cycle()  # Cycle to the next content everytime we call the model
return action

model = CyclerModel(input_dim=2, n_arms=3)

input_arr = tf.random.uniform((2, ))
outputs = model._set_inputs(input_arr)

# keras.saving.save_model(model, 'tmp')
# model = keras.saving.load_model('tmp')

inputs = np.array([[1.0, 2.0], [1.0, 2.0], [1.1, 2.0], [1.1, 2.0], [1.0, 2.0]], dtype=np.float32)
outputs = [model(arr[None]) for arr in inputs]
print(outputs)  # expected output 0, 1, 2, 0, 1

Приведенный выше фрагмент кода отлично работает, когда строки сохранения комментируются, что дает ожидаемый результат.
Точная ошибка выглядит следующим образом:< /p>

Traceback (последний вызов — последний):
Файл «xxx», строка 63, в
keras.saving.save_model (модель, 'tmp')
Файл "xxx", строка 167, в save_model
return Legacy_sm_saving_lib.save_model(
^^^^^^^^^^^^^^^^^^^^^^^^^^^ ^^^^
Файл "xxx/keras/src/utils/traceback_utils.py", строка 70, в error_handler
raise e.with_traceback(filtered_tb) из None
Файл "xxx/tensorflow/python/ops/op_selector.py", строка 417, в map_subgraph
raise UnliftableError(
tensorflow.python.ops.op_selector .UnliftableError: невозможно поднять тензор , поскольку он транзитивно зависит от заполнителя через at хотя бы один путь, например: Cycler/Variable/Initializer/initial_value (Pack)

Подробнее здесь: https://stackoverflow.com/questions/793 ... tensor-bec