Element-wise multiplication with Keras

Question

I have a RGB image of shape (256,256,3) and I have a weight mask of shape (256,256). How do I perform the element-wise multiplication between them with Keras? (all channels share the same mask)

Answer 1

You need a Reshape so both tensors have the same number of dimensions, and a Multiply layer

mask = Reshape((256,256,1))(mask) 
out = Multiply()([image,mask])

If you have variable shapes, you can use a single Lambda layer like this:

import keras.backend as K 

def multiply(x):
    image,mask = x
    mask = K.expand_dims(mask, axis=-1) #could be K.stack([mask]*3, axis=-1) too 
    return mask*image

out = Lambda(multiply)([image,mask])

Answer 2

As an alternative you can do this using a Lambda layer (as in @DanielMöller's answer you need to add a third axis to the mask):

from keras import backend as K

out = Lambda(lambda x: x[0] * K.expand_dims(x[1], axis=-1))([image, mask])