Hello everyone,
I am currently trying to create a tf.function for my tensorflow lite model that takes a dictionary of weights for different tensors and loads it into the lite interpreter. Currently it is possible to use my function by using a workaround where I hardcode the input signature for the different dictionary values, but I want to know if there is a better solution.
The dictionary I want to pass looks something like this (In my code i generated the dictionary, this here is just to demonstrate the form of the dictionary:
weights_dict = {
"conv1d/bias:0": [0., 0., 0., 0., 0.],
"conv1d/kernel:0": [[[-0.23031473, 0.25823927, -0.15676963, -0.27285957, 0.36472344]],
[[-0.15187979, -0.00110114, 0.05242634, 0.2745613, -0.20043385]],
[[ 0.03774023 -0.45265818 0.24690723 0.23345459 0.15752888]],
[[-0.3444885 -0.07790518 -0.11909294 -0.34963953 0.14239442]]],
"dense/bias:0": [0. 0. 0. 0. 0. 0. 0. 0.],
"dense/kernel:0": [[ 0.10110682, -0.50597924, 0.3949796, 0.20653826, -0.26708043, 0.2091226, -0.5352592, -0.6049661],
[-0.4578067, 0.29674542, -0.27303556, -0.568668, 0.4862703, 0.5257646, -0.28185177, 0.43704462],
[-0.45960283, -0.44893104, -0.386113, -0.18281126, -0.44102007, -0.21802643, 0.60579884, 0.23984993],
[ 0.45038152, 0.5851245, 0.40392494, 0.28304034, 0.6369395, -0.06007874, -0.5181305, -0.6644601],
[-0.42245385, 0.43564594, 0.45273066, 0.5162871, 0.12010175, -0.4153296, -0.00371552, 0.17521149]]
"dense_1/bias:0": [0.],
"dense_1/kernel:0": [[0.18899703],
[0.54743993],
[0.57986426],
[-0.31321746],
[0.37268388],
[0.18092150],
[-0.14566028],
[-0.26471186]]
}
This dictionary is passed into this tf.function:
@tf.function(input_signature=[signature_dict])
def set_weights(self, weights):
tf.print("im in")
tensor_names = [weight.name for weight in self.model.weights]
for i, tensor in enumerate(self.model.weights):
tensor.assign(weights[tensor_names[i]])
return tensor_names
The workaround i mentioned before is the signature_dict i pass into the input_signature. It is a hardcoded representation of the tensor shapes and looks like this:
signature_dict = { "conv1d/bias:0": tf.TensorSpec(shape=[5], dtype=tf.float32),
"conv1d/kernel:0": tf.TensorSpec(shape=[4, 1, 5], dtype=tf.float32),
"dense/bias:0": tf.TensorSpec(shape=[8], dtype=tf.float32),
"dense/kernel:0": tf.TensorSpec(shape=[5, 8], dtype=tf.float32),
"dense_1/bias:0": tf.TensorSpec(shape=[1], dtype=tf.float32),
"dense_1/kernel:0": tf.TensorSpec(shape=[8, 1], dtype=tf.float32) }
In order to invoke this function I need to do this:
set_weights = interpreter.get_signature_runner("set_weights")
weights = weights_dict["conv1d/bias:0"]
weights_1 = weights_dict["conv1d/kernel:0"]
weights_2 = weights_dict["dense/bias:0"]
weights_3 = weights_dict["dense/kernel:0"]
weights_4 = weights_dict["dense_1/bias:0"]
weights_5 = weights_dict["dense_1/kernel:0"]
output= set_weights(weights=weights, weights_1=weights_1, weights_2=weights_2, weights_3=weights_3, weights_4=weights_4, weights_5=weights_5)
As you can see, i have to do this weird workaround where I have to manually create the signature TensorSpecs and extract the different weight tensors from the dictionary before passing them using 6 different variables.
I would really like to simplify this, because it does not look good and I have to change everything if I want to change the model layers. In the best case, I want to pass the dictionary as a whole to the set_weights function and let it do the rest. This should like this:
output= set_weights(weights=weight_dict)
Does anybody know if there is a way to pass the dictionary as a whole into the tf.function and how? Thanks in advance!