Блог

Я хочу расширить tf.keras.Model для реализации обратимых нейронных сетей и сохранить / загрузить эту модель с помощью tf.keras.Model.save() amp; tf.keras.models.load_model()

В частности, я написал два класса:

• Обратимый класс NN
  • подклассы tf.keras.Model
  • имеет inverse() метод
• Класс связующего слоя
  • подклассы tf.keras.layers.Layer
  • имеет inverse() метод
  • имеет log_det_jacobian свойство ( @property ).

Код, связанный с этим, выглядит следующим образом (пожалуйста, обратитесь к моему репозиторию github для получения полного кода).

 class SimpleRealNVP(K.Model):
    ...
    def inverse(self, inputs):
        h = inputs
        for layer in self.couplings:
            h = self.layer.inverse(h)
        return h
    ...
  

 class AffineCoupling(L.Layer):
    ....
    @property
    def log_det_jacobian(self):
        return self._log_det_jacobian

    def inverse(self, input):
        ...
  

Я просто попробовал tf.keras.Model.save() и tf.keras.load_model() , но я получил сообщение об ошибке, в котором говорилось

 ...
mdl = K.models.load_model("path/to/savedmodel")
...
x_manifold = mdl.inverse(z_manifold)

# Traceback (most recent call last):
#   File ".manifold.py", line 54, in <module>
#     x_manifold = mdl.inverse(z_manifold)
# AttributeError: 'SimpleRealNVP' object has no attribute 'inverse'
  

Какой подходящий способ включить расширенные методы / свойства в SavedModel?

(Я также пытался указать custom_objects in load_model() , но это не сработало.)