Keras layer channel-wise multiplication of scalar and graph plotting - tensorflow

Question

I try to multiply scalar values to each channel in a tensor:
import tensorflow as tf
t = tf.ones([2,3,3,4])
w = tf.constant([1,2,3,4], dtype=tf.float32)
tf.multiply(t,w)
yields
<tf.Tensor: shape=(2, 3, 3, 4), dtype=float32, numpy=
array([[[[1., 2., 3., 4.],
[1., 2., 3., 4.],
[1., 2., 3., 4.]],
...
which is correct.
Now I am trying to wrap that operation inside a keras.layers.Layer, whereby w is a learnable parameter. I also try to plot my model using tf.keras.utils.plot_model(m). I encounter several problems.
Method 1
from tensorflow.keras import Model, Input
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.w = self.add_weight(shape=(256,), trainable=True)
def call(self, x):
return x * self.w
I plot this model using
mm = MyModel()
x = Input(shape=(64, 64, 256), batch_size=10, name='Input')
m = Model(inputs=[x], outputs=mm.call(x))
tf.keras.utils.plot_model(m)
Problem: I encountered the following warning:
WARNING:tensorflow:
The following Variables were used a Lambda layer's call (tf.math.multiply_2), but
are not present in its tracked objects:
<tf.Variable 'Variable:0' shape=(256,) dtype=float32>
It is possible that this is intended behavior, but it is more likely
an omission. This is a strong indication that this layer should be
formulated as a subclassed Layer rather than a Lambda layer.
Question: Can I savely ignore the warning and the weights are still learned? If yes, how can I suppress this warning?
Method 2
As suggested in the warning, I wrap the multiplication in its own subclassed layer:
class MyMultiply(Layer):
def __init__(self):
super(MyMultiply, self).__init__()
def call(self, x):
return tf.multiply(x[0], x[1])
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.w = self.add_weight(shape=(256,), trainable=True)
self.mul = MyMultiply()
def call(self, x):
return self.mul([x, self.w])
Problem: This works until the model is plotted. Then I encounter the following error: AttributeError: 'ResourceVariable' object has no attribute '_keras_history'
Traceback:
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
<ipython-input-18-e4cc5cc97726> in <module>()
21 x = Input(shape=(64, 64, 256), batch_size=10, name='Input')
22 m = Model(inputs=[x], outputs=mm.call(x))
---> 23 tf.keras.utils.plot_model(m)
---------------------------------------------------------------------------
/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/node.py in <lambda>(t)
259 if self.is_input:
260 return []
--> 261 inbound_layers = nest.map_structure(lambda t: t._keras_history.layer,
262 self.call_args[0])
263 return inbound_layers
AttributeError: 'ResourceVariable' object has no attribute '_keras_history'
Question: How do I resolve that error? Is this a bug (I submitted an issue to the tf github repo, however it was deleted immediately)?
Method 3
I try to use keras.layers.Multiply instead:
from tensorflow.keras.layers import Multiply
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.w = self.add_weight(shape=(256,), trainable=True)
self.mul = Multiply()
def call(self, x):
return self.mul([x, self.w])
Problem: ValueError: Can not merge tensors with different batch sizes. Got tensors with shapes : [(10, 64, 64, 256), (256,)]
To my understanding, the ValueError occurs because the internal _Merge layer checks for equal batch sizes. The internal Multiply layer however implements the multiplication with broadcasting (which should work!):
# from tensorflow/python/keras/layers/merge.py line 316-320
def _merge_function(self, inputs):
output = inputs[0]
for i in range(1, len(inputs)):
output = output * inputs[i]
return output
I could use tf.broadcast_to and so on, however, to my understanding this would materialize the tensor and would occupy more memory which I try to avoid.
Question: Is there another way to make keras.layers.Multiply work, so ultimately the model plotting works?

Answer 1

You can avoid the warning in Method 1 by creating a Keras Layer instead of a Model.
import tensorflow as tf
class MyLayer(tf.keras.layers.Layer):
def __init__(self):
super(MyLayer, self).__init__()
self.w = self.add_weight(name='multiply_weight', shape=(256,), trainable=True)
def call(self, x):
return tf.multiply(x, self.w)
mul_layer = MyLayer()
x = tf.keras.Input(shape=(64, 64, 256), batch_size=10, name='Input')
output = mul_layer(x)
m = tf.keras.Model(inputs=[x], outputs=output)
tf.keras.utils.plot_model(m)