How to set weights in model signature?

Hello everyone!

I deal with federated learning on device and because so, my Android code needs to set the weights of a model in TFLite format. I’m trying to create a model signature to set the weights, without success.

Following are my trials and the error messages for each one.

Please, anyone could help me? Thanks in advance!

# pip show tensorflow
Name: tensorflow
Version: 2.11.0
def convert_and_save(model, saved_model_dir):
    tf.saved_model.save(
        model,
        saved_model_dir,
        signatures=model.add_signatures(),
        )
    
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
        tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.
    ]
    converter.experimental_enable_resource_variables = True
    tflite_model = converter.convert()
  
    model_file_path = os.path.join(saved_model_dir, "model.tflite")
    with open(model_file_path, "wb") as model_file:
        model_file.write(tflite_model)

Trial 1

    def add_signatures(self):
      @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32,)])
      def initialize_weights(weights):
          count = 0
          for var in self.model.trainable_weights:
            var.assign(weights[count])
            count += 1

ValueError: Cannot assign value to variable ’ head_model_dense_1/kernel:0’: Shape mismatch.The variable shape (24, 16), and the assigned value shape () are incompatible.

Trial 2

    def add_signatures(self):
      @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32,)])
      def initialize_weights(weights):
          self.model.set_weights(weights)

File “/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py”, line 1783, in set_weights **
if expected_num_weights != len(weights):

TypeError: len is not well defined for a symbolic Tensor (weights:0). Please call `x.shape` rather than `len(x)` for shape information.

Trial 3

ws0 = tf.TensorSpec(shape=(NUM_BOTTLENECK_FEATURES, N_NEURONS_HEAD_MODEL_LAYER), dtype=tf.float32, name='head_model_dense_1/kernel:0')
ws1 = tf.TensorSpec(shape=(N_NEURONS_HEAD_MODEL_LAYER,), dtype=tf.float32, name="head_model_dense_1/bias:0")
ws2 = tf.TensorSpec(shape=(N_NEURONS_HEAD_MODEL_LAYER, NUM_CLASSES), dtype=tf.float32, name="head_model_softmax/kernel:0")
ws3 = tf.TensorSpec(shape=(NUM_CLASSES,), dtype=tf.float32, name="head_model_softmax/bias:0")

weights_signature_array = [ws0, ws1, ws2, ws3]

    def add_signatures(self):
      @tf.function(input_signature=weights_signature_array)
      def initialize_weights(w0,w1,w2,w3):
          restored_tensors = {}
          weights = [w0,w1,w2,w3]
          count = 0
          for var in self.head_model.trainable_weights:
            var.assign(weights[count])
            restored_tensors[var.name] = weights[count]
            count += 1

          return restored_tensors

File “/usr/local/lib/python3.8/dist-packages/tensorflow/python/saved_model/save.py”, line 549, in _map_function_arguments_to_created_inputs
raise ValueError(
ValueError: Got non-flat/non-unique argument names for SavedModel signature ‘initialize’: more than one argument to ‘__inference_signature_wrapper_43205’ was named ‘head_model_dense_1/bias:0’. Signatures have one Tensor per named input, so to have predictable names Python functions used to generate these signatures should avoid *args and Tensors in nested structures unless unique names are specified for each. Use tf.TensorSpec(…, name=…) to provide a name for a Tensor input.

Hi @Alexandre_Osorio, generally Signatures in TensorFlow Lite are used to

  • Specify inputs and outputs of the converted TensorFlow Lite model by respecting the TensorFlow model’s signatures.
  • Allow a single TensorFlow Lite model to support multiple entry points.

I don’t think you can be able to pass model weights to model signatures. Thank You.