def test_tensorflow_wrapper_keras_subclass_decorator_compile_args(): import tensorflow as tf class UndecoratedModel(tf.keras.Model): def call(self, inputs): return inputs # Can't wrap an undecorated keras subclass model with pytest.raises(ValueError): TensorFlowWrapper(UndecoratedModel()) @keras_subclass( "TestModel", X=numpy.array([0.0, 0.0]), Y=numpy.array([0.5]), input_shape=(2,), compile_args={"loss": "binary_crossentropy"}, ) class TestModel(tf.keras.Model): def call(self, inputs): return inputs model = TensorFlowWrapper(TestModel()) model = model.from_bytes(model.to_bytes()) assert model.shims[0]._model.loss == "binary_crossentropy" assert isinstance(model, Model)
def test_tensorflow_wrapper_keras_subclass_decorator_capture_args_kwargs( X, Y, input_size, n_classes, answer ): import tensorflow as tf @keras_subclass( "TestModel", X=numpy.array([0.0, 0.0]), Y=numpy.array([0.5]), input_shape=(2,) ) class TestModel(tf.keras.Model): def __init__(self, custom=False, **kwargs): super().__init__(self) # This is to force the mode to pass the captured arguments # or fail. assert custom is True assert kwargs.get("other", None) is not None def call(self, inputs): return inputs # Can wrap an decorated keras subclass model model = TensorFlowWrapper(TestModel(True, other=1337)) assert hasattr(model.shims[0]._model, "eg_args") args_kwargs = model.shims[0]._model.eg_args assert True in args_kwargs.args assert "other" in args_kwargs.kwargs # Raises an error if the args/kwargs is not serializable obj = {} obj["key"] = obj with pytest.raises(ValueError): TensorFlowWrapper(TestModel(True, other=obj)) # Provides the same arguments when copying a capture model model = model.from_bytes(model.to_bytes())
def test_tensorflow_wrapper_serialize_model_subclass( X, Y, input_size, n_classes, answer ): import tensorflow as tf input_shape = (1, input_size) ops = get_current_ops() @keras_subclass( "foo.v1", X=ops.alloc2f(*input_shape), Y=to_categorical(ops.asarray1i([1]), n_classes=n_classes), input_shape=input_shape, ) class CustomKerasModel(tf.keras.Model): def __init__(self, **kwargs): super(CustomKerasModel, self).__init__(**kwargs) self.in_dense = tf.keras.layers.Dense( 12, name="in_dense", input_shape=input_shape ) self.out_dense = tf.keras.layers.Dense( n_classes, name="out_dense", activation="softmax" ) def call(self, inputs) -> tf.Tensor: x = self.in_dense(inputs) return self.out_dense(x) model = TensorFlowWrapper(CustomKerasModel()) # Train the model to predict the right single answer optimizer = Adam() for i in range(50): guesses, backprop = model(X, is_train=True) d_guesses = (guesses - Y) / guesses.shape[0] backprop(d_guesses) model.finish_update(optimizer) predicted = model.predict(X).argmax() assert predicted == answer # Save then Load the model from bytes model.from_bytes(model.to_bytes()) # The from_bytes model gets the same answer assert model.predict(X).argmax() == answer