I defined simple custom model:
import tensorflow as tf
class CustomModule(tf.keras.layers.Layer):
def __init__(self):
super(CustomModule, self).__init__()
self.v = tf.Variable(1.)
def call(self, x):
print('Tracing with', x)
return x * self.v
def mutate(self, new_v):
self.v.assign(new_v)
I want to save it for serving and that is why I need to provide a function for “serving_default”. I’ve tried to do it like this:
module = CustomModule() module_with_signature_path = './tmp/1' call = tf.function(module.mutate, input_signature=[tf.TensorSpec([], tf.float32)]) tf.saved_model.save(module, module_with_signature_path, signatures=call)
I got an error:
ValueError: Got a non-Tensor value <tf.Operation 'StatefulPartitionedCall' type=StatefulPartitionedCall> for key 'output_0' in the output of the function __inference_mutate_8 used to generate the SavedModel signature 'serving_default'. Outputs for functions used as signatures must be a single Tensor, a sequence of Tensors, or a dictionary from string to Tensor.
How can I properly define signature while saving model? Thank you!
Hi all: this is one of my first posts on Stackoverflow - so apologies in advance if i'm not conforming to certain standards!
I'm having trouble saving my Keras model as a mlflow.pyfunc model as it's giving me a "cannot pickle a 'weakref' object when I try to log it.
So why am i saving my Keras model as a pyfunc model object in the first place? This is because I want to override the default predict method and output something custom. I also want to do some pre-processing steps on the X_test or new data by encoding it with a tf.keras.StringLookup and then invert it back to get the original categorical variable class. For this reason, I was advised by Databricks that the mlflow.pyfunc flavor is the best way to go for these types of use-cases
The Keras model works just fine and i'm able to log it using mlflow.keras.log_model. But it fails when i try to wrap it inside a cutomer "KerasWrapper" class.
Here are some snippets of my code. For the purpose of debugging, the current predict method in the custom class is just the default. I simplified it to help debug, but obviously I haven't been able to resolve it.
I would be extremely grateful for any help. Thanks in advance!
ALL CODE ON AZURE DATABRICKS
Custom mlflow.pyfunc class
class KerasWrapper(mlflow.pyfunc.PythonModel):
def __init__(self, keras_model, labelEncoder, labelDecoder, n):
self.keras_model = keras_model
self.labelEncoder = labelEncoder
self.labelDecoder = labelDecoder
self.topn = n
def load_context(self, context):
self.keras_model = mlflow.keras.load_model(model_uri=context.artifacts[self.keras_model], compile=False)
def predict(self, context, input_data):
scores = self.keras_model.predict(input_data)
return scores
My Keras Deep Learning Model (this works fine by the way)
def build_model(vocab_size, steps, drop_embed, n_dim, encoder, modelType):
model = None
i = Input(shape=(None,), dtype="int64")
#embedding layer
e = Embedding(vocab_size, 16)(i)
s = SpatialDropout1D(drop_embed)(e)
x = Conv1D(256, steps, activation='relu')(s)
x = GlobalMaxPooling1D()(x)
x = Dense(128, activation='relu')(x)
x = Dropout(0.2)(x)
#output layer
x = Dense(vocab_size, activation='softmax')(x)
model = Model(i, x)
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.summary()
return model
MLFLOW Section
with mlflow.start_run(run_name=runName):
mlflow.tensorflow.autolog()
#Build the model, compile and train on the training set
#signature: build_model(vocab_size, steps, drop_embed, n_dim, encoder, modelType):
keras_model = build_model((vocab_size + 1), timeSteps, drop_embed, embedding_dimensions, encoder, modelType)
keras_model.fit(X_train_encoded, y_train_encoded, epochs=epochs, verbose=1, batch_size=32, use_multiprocessing = True,
validation_data=(X_test_encoded, y_test_encoded))
# Log the model parameters used for this run.
mlflow.log_param("numofActionsinWorkflow", numofActionsinWf)
mlflow.log_param("timeSteps", timeSteps)
#wrap it up in a pyfunc model
wrappedModel = KerasWrapper(keras_model, encoder, decoder, bestActionCount)
# Create a model signature using the tensor input to store in the MLflow model registry
signature = infer_signature(X_test_encoded, wrappedModel.predict(None, X_test_encoded))
# Let's check out how it looks
print(signature)
# Create an input example to store in the MLflow model registry
input_example = np.expand_dims(X_train[17], axis=0)
# The necessary dependencies are added to a conda.yaml file which is logged along with the model.
model_env = mlflow.pyfunc.get_default_conda_env()
# Record specific additional dependencies required by the serving model
model_env['dependencies'][-1]['pip'] += [
f'tensorflow=={tf.__version__}',
f'mlflow=={mlflow.__version__}',
f'sklearn=={sklearn.__version__}',
f'cloudpickle=={cloudpickle.__version__}',
]
#log the model to experiment
#mlflow.keras.log_model(keras_model, artifact_path = runName, signature=signature, input_example=input_example, conda_env = model_env)
wrapped_model_path = runName
if (os.path.exists(wrapped_model_path)):
shutil.rmtree(wrapped_model_path)
#Log model as pyfunc model
mlflow.pyfunc.log_model(runName, python_model=wrappedModel, signature=signature, input_example=input_example, conda_env = model_env)
#return the run ID for model registration
run_id = mlflow.active_run().info.run_id
mlflow.end_run()
Here is the error that i receive
I am getting an error, when I run the below code.
def __init__(self, model, sess, loss_fn=None):
"""
To generate the White-box Attack Agent.
:param model: the target model which should have the input tensor, the target tensor and the loss tensor.
:param sess: the tensorflow session.
:param loss_fn: None if using original loss of the model.
"""
self.model = model
self.input_tensor = model.inputs[0]
self.output_tensor = model.outputs[0]
self.target_tensor = model.targets[0]
self._sample_weights = model.sample_weights[0]
if loss_fn is None:
self.loss_tensor = model.total_loss
self.gradient_tensor = K.gradients(self.loss_tensor, self.input_tensor)[0]
else:
self.set_loss_function(loss_fn)
self.sess = sess
error:
self.target_tensor = model.targets[0] , AttributeError: 'Model' object has no attribute 'targets'
I am working with Tensorflow 1.14.0 ,keras 2.2.4-tf and python 3.6.13.how can I resolve this problem?
Thank you
The model argument that your function is using is apparently talking about a unique class that isn't part of the default Keras or Tensorflow. You need to make sure that the model object you're passing the function is an instance of the appropriate class. Or, alternatively, you can jerry rig it by doing something like this:
# whatever code instantiates the `model` would be here
...
...
model.targets = [ "targets of an appropriate structure and type go here" ]
Then when you pass that model object to the function, it will have a targets attribute to be accessed. But it's better to just use the right object in the first place.
Or, the third option: just comment out the offending line and hope it isn't actually used anywhere (if it's being set, it's probably used somewhere or other though).
def __init__(self, model, sess, loss_fn=None):
"""
To generate the White-box Attack Agent.
:param model: the target model which should have the input tensor, the target tensor and the loss tensor.
:param sess: the tensorflow session.
:param loss_fn: None if using original loss of the model.
"""
self.model = model
self.input_tensor = model.inputs[0]
self.output_tensor = model.outputs[0]
#self.target_tensor = model.targets[0]
self._sample_weights = model.sample_weights[0]
if loss_fn is None:
self.loss_tensor = model.total_loss
self.gradient_tensor = K.gradients(self.loss_tensor, self.input_tensor)[0]
else:
self.set_loss_function(loss_fn)
self.sess = sess
I have wrote a custom code to build a UNet architecture. To do so I have firstly subclassed the tf.keras.layers.Layer object to define an encoder convolutional block composed by a conv3D layer, a BatchNormalization layer and a Activation layer, similarly I defined a decoder inverse convolutional block composed by a Conv3DTranspose layer, a BatchNormalization layer, an Activation layer and a Concatenate layer. Finally I subclassed the tf.keras.Model object to define the full model, composed by 4 enconding blocks and 4 decoding blocks.
To checkpoint the model while training I have used the tf.keras.callbacks.ModelCheckpoint callback. However when a I try to load back the model (that in fact is still training) with tf.keras.models.load_model() I receive the following error: ValueError: No model found in config file.
Here the full code for the model definition, building and fitting:
import tensorflow as tf
# Encoder block
class ConvBlock(tf.keras.layers.Layer):
def __init__(self, n_filters, conv_size, conv_stride, **kwargs):
super(ConvBlock, self).__init__(**kwargs)
self.conv3D = tf.keras.layers.Conv3D(
filters=n_filters,
kernel_size=conv_size,
strides=conv_stride,
padding="same",
)
self.batch_norm = tf.keras.layers.BatchNormalization()
self.relu = tf.keras.layers.Activation("relu")
def call(self, inputs, training=None):
h = self.conv3D(inputs)
if training:
h = self.batch_norm(h)
h = self.relu(h)
return h
# Decoder block
class InvConvBlock(tf.keras.layers.Layer):
def __init__(self, n_filters, conv_size, conv_stride, activation, **kwargs):
super(InvConvBlock, self).__init__(**kwargs)
self.conv3D_T = tf.keras.layers.Conv3DTranspose(
filters=n_filters,
kernel_size=conv_size,
strides=conv_stride,
padding="same",
)
self.batch_norm = tf.keras.layers.BatchNormalization()
self.activ = tf.keras.layers.Activation(activation)
self.concat = tf.keras.layers.Concatenate(axis=-1)
def call(self, inputs, feat_concat=None, training=None):
h = self.conv3D_T(inputs)
if training:
h = self.batch_norm(h)
h = self.activ(h)
if feat_concat is not None:
h = self.concat([h, feat_concat])
return h
class UNet(tf.keras.Model):
def __init__(self, n_filters, e_size, e_stride, d_size, d_stride, **kwargs):
super(UNet, self).__init__(**kwargs)
# Encoder
self.conv_block_1 = ConvBlock(n_filters, e_size, e_stride)
self.conv_block_2 = ConvBlock(n_filters * 2, e_size, e_stride)
self.conv_block_3 = ConvBlock(n_filters * 4, e_size, (1, 1, 1))
self.conv_block_4 = ConvBlock(n_filters * 8, e_size, (1, 1, 1))
# Decoder
self.inv_conv_block_1 = InvConvBlock(n_filters * 4, d_size, (1, 1, 1), "relu")
self.inv_conv_block_2 = InvConvBlock(n_filters * 2, d_size, (1, 1, 1), "relu")
self.inv_conv_block_3 = InvConvBlock(n_filters, d_size, d_stride, "relu")
self.inv_conv_block_4 = InvConvBlock(1, d_size, d_stride, "sigmoid")
def call(self, inputs, **kwargs):
h1 = self.conv_block_1(inputs, **kwargs)
h2 = self.conv_block_2(h1, **kwargs)
h3 = self.conv_block_3(h2, **kwargs)
h = self.conv_block_4(h3, **kwargs)
h = self.inv_conv_block_1(h, feat_concat=h3, **kwargs)
h = self.inv_conv_block_2(h, feat_concat=h2, **kwargs)
h = self.inv_conv_block_3(h, feat_concat=h1, **kwargs)
h = self.inv_conv_block_4(h, **kwargs)
return h
model = UNet(
n_filters,
e_size,
e_stride,
d_size,
d_stride,
)
model.build((None, *input_shape, 1))
loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam(learning_rate)
metrics = [tf.keras.metrics.Precision(), tf.keras.metrics.Recall()]
model.compile(
loss=loss,
optimizer=optimizer,
metrics=metrics,
)
CP_callback = tf.keras.callbacks.ModelCheckpoint(
f"{checkpoint_dir}/model.h5", save_freq='epoch', monitor="loss"
)
unet.fit(
data,
epochs=opts.epochs,
callbacks=[CP_callback],
)
To load the model I used the following code on another python console:
import tensorflow as tf
model = tf.keras.models.load_model(f'{checkpoint_dir}/model.h5')
but here I receive the above mentioned error. What am I missing? Or what am I doing wrong?
Thank you in advance for your help.
This is because you don't define the get_config method in your custom layers. For this check, this exited answer in SO.
Otherwise, you can save the trained weights (not the full model) and load the model as follows. In that case, you don't need to define this function. Please note, it's good practice to do, however. Here is a workaround for your problem:
# callback
tf.keras.callbacks.ModelCheckpoint('model.h5',
monitor='val_loss',
verbose= 1,
save_best_only=True,
mode= 'min',
save_weights_only=True) # <---- only save weight
# train
model = UNet(
n_filters,
e_size,
e_stride,
d_size,
d_stride,
)
model.compile(...)
model.fit(...)
# inference
model = UNet(
n_filters,
e_size,
e_stride,
d_size,
d_stride,
)
model.build((None, *input_shape, 1))
model.load_weights('model.h5')
For more details, see the documentation of Serialization and saving and also collab demonstration of François Chollet. Also, We've written an article about model subclassing and custom training stuff in tf 2.x, in the Save and Load section (at the bottom) of this article, we've demonstrated many strategies, here, hope that help.
Update
I've run your public colab notebook. Unfortunately, I am facing the same issue, and it's a bit weird and currently, I don't have the exact answer for saving the entire model in the ModelCheckpoint callback with Custom Layer even if we define the get_config() method.
However, there is another workaround that may come in handy for you. As we know there are two major ways to save tf models: (1). SaveModel and HDF5 format. The way is we choose the SaveMoedl format. Which is recommended by the way and safe to use.
The key difference between HDF5 and SavedModel is that HDF5 uses object configs to save the model architecture, while SavedModel saves the execution graph. Thus, SavedModels are able to save custom objects like subclassed models and custom layers without requiring the original code.
Now, as for your requirements, you are saving the entire model along with the best loss or val_loss in training time. For that, we can define a custom callback do save the model for lowest validation_loss (or whatever you want). As follows:
class SaveModelH5(tf.keras.callbacks.Callback):
def on_train_begin(self, logs=None):
self.val_loss = []
def on_epoch_end(self, epoch, logs=None):
current_val_loss = logs.get("val_loss")
self.val_loss.append(logs.get("val_loss"))
if current_val_loss <= min(self.val_loss):
print('Find lowest val_loss. Saving entire model.')
self.model.save('unet', save_format='tf') # < ----- Here
save_model = SaveModelH5()
unet.fit(.., callbacks=save_model)
Using
model.save('any_name', save_format=`tf`)
allows us create a any_name working directory, inside which it contains assets, saved_model.pb, and variables. The model architecture and training configuration, including the optimizer, losses, and metrics are stored in saved_model.pb. The weights are saved in the variables directory.
When saving the model and its layers, the SavedModel format stores the class name, call function, losses, and weights (and the config, if implemented). The call function defines the computation graph of the model/layer. In the absence of the model/layer config, the call function is used to create a model that exists like the original model which can be trained, evaluated, and used for inference. When we need to re-load the saved model, we can do as follows:
new_unet = tf.keras.models.load_model("unet", compile=False)
Colab.
I was trying to implement the Google Bert model in tensorflow-keras using tensorflow hub. For this I designed a custom keras layer "Bertlayer" . Now the problem is when I am compiling the keras model it keeps showing that
AttributeError: 'Bertlayer' object has no attribute '_keras_style'
I don't know where I am wrong and what _keras_style attribute is.Please help to find the error in the code.
This is the github link to the full code: https://github.com/PradyumnaGupta/BERT/blob/master/Untitled21.ipynb
class BertLayer(tf.layers.Layer):
def __init__(self, n_fine_tune_layers=10, **kwargs):
self.n_fine_tune_layers = n_fine_tune_layers
self.trainable = True
self.output_size = 768
super(BertLayer, self).__init__(**kwargs)
def build(self, input_shape):
self.bert = hub.Module(
bert_path,
trainable=self.trainable,
name="{}_module".format(self.name)
)
trainable_vars = self.bert.variables
# Remove unused layers
trainable_vars = [var for var in trainable_vars if not "/cls/" in var.name]
# Select how many layers to fine tune
trainable_vars = trainable_vars[-self.n_fine_tune_layers :]
# Add to trainable weights
for var in trainable_vars:
self._trainable_weights.append(var)
for var in self.bert.variables:
if var not in self._trainable_weights:
self._non_trainable_weights.append(var)
super(BertLayer, self).build(input_shape)
def call(self, inputs):
inputs = [K.cast(x, dtype="int32") for x in inputs]
input_ids, input_mask, segment_ids = inputs
bert_inputs = dict(
input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids
)
result = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
"pooled_output"
]
return result
def compute_output_shape(self, input_shape):
return (input_shape[0], self.output_size)
So, tensorflow version 1.* is a bit misleading. It actually has 2 bases classes called Layer. One - the one that you are using. It is intended to implement shortcut wrappers over regular TF operations. The other from tensorflow.keras.layers import Layer is for Keras-like models and sequencies.
Judging by your error, you are using keras/models to train further.
You probably should start form derivering your layer from keras.layers.Layer instead of tf.layers.Layer.