Keras: Matrix size-incompatible: In[0]: [1,16384], In[1]: [1024,10]

Question

I am trying to implement gaussian attention with keras+tensorflow like it is described here: http://akosiorek.github.io/ml/2017/10/14/visual-attention.html#mjx-eqn-att

For this, I wrote a custom Keras layer like this (I adjusted the gaussian_mask method a little bit compared to the blog post):

def gaussian_mask(u, s, d, R, C, transpose=False):
    """
    :param u: tf.Tensor, centre of the first Gaussian.
    :param s: tf.Tensor, standard deviation of Gaussians.
    :param d: tf.Tensor, shift between Gaussian centres.
    :param R: int, number of rows in the mask, there is one Gaussian per row.
    :param C: int, number of columns in the mask.
    """
    # indices to create centres
    R = tf.to_float(tf.reshape(tf.range(R), (R, 1, 1)))
    C = tf.to_float(tf.reshape(tf.range(C), (1, C, 1)))


    centres = u[:, np.newaxis, np.newaxis] + R * d
    column_centres = C - centres
    mask = tf.exp(-.5 * tf.square(column_centres / s))
    # we add eps for numerical stability
    normalised_mask = mask / (tf.reduce_sum(mask, 1, keep_dims=True) + 1e-8)

    return normalised_mask

class visual_attention_layer(Layer):
    def __init__(self, output_dim, transpose=False, **kwargs):
        self.output_dim = output_dim
        self.transpose = transpose
        super(visual_attention_layer, self).__init__(**kwargs)

    def build(self, input_shape):
        super(visual_attention_layer, self).build(input_shape)

    def call(self, x): 
        x_x, x_y, input_img = x

        u_x,s_x,d_x = tf.split(x1,3,1)
        u_y,s_y,d_y = tf.split(x2,3,1)


        W = input_img.shape[1]
        H = W = input_img.shape[2]
        Ay = gaussian_mask(u_y, s_y, d_y, self.output_dim[0], H)
        Ax = gaussian_mask(u_x, s_x, d_x, self.output_dim[0], W)

        input_img = tf.transpose(input_img, perm=[0,3,1,2])
        Ay = tf.transpose(Ay, perm=[0, 3, 1, 2])
        Ax = tf.transpose(Ax, perm=[0, 3, 1, 2])


        glimpse = tf.matmul( input_img, Ax, transpose_b=True)
        glimpse = tf.matmul(Ay, glimpse)
        glimpse = tf.transpose(glimpse, perm=[0,2,3,1])

        return glimpse

    def compute_output_shape(self, input_shape):
        return (self.output_dim[0], self.output_dim[1], input_shape[2][3])

and then use it like this:

inputs = Input(shape=(28,28,1))

x = Conv2D(64, kernel_size=(3,3), activation="relu")(inputs)
x = MaxPool2D()(x)
x = Conv2D(64, kernel_size=(3,3), activation="relu")(x)
x = MaxPool2D()(x)
x = Flatten()(x)

x1 = Dense(3, activation="sigmoid")(x)
x2 = Dense(3, activation="sigmoid")(x)
x = visual_attention_layer(output_dim=(20,20))([x1,x2, inputs])

x = Conv2D(64, kernel_size=(3,3), activation="relu")(x)
#x = MaxPool2D()(x)
x = Conv2D(64, kernel_size=(3,3), activation="relu")(x)
x = Flatten()(x)
predictions = Dense(10, activation='softmax')(x)

model = Model(inputs=inputs, outputs=predictions)
model.compile(optimizer='rmsprop',
              loss='categorical_crossentropy',
              metrics=['accuracy'])
history = model.fit(x_train, y_train, epochs=5, batch_size=1)  

The model compiles fine (except when I use the MaxPool2D that is commented out right now, than I get a "IndexError: tuple index out of range"). However, when I want to train it, I get the following error:

InvalidArgumentError                      Traceback (most recent call last)
C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\client\session.py in _do_call(self, fn, *args)
   1321     try:
-> 1322       return fn(*args)
   1323     except errors.OpError as e:

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\client\session.py in _run_fn(feed_dict, fetch_list, target_list, options, run_metadata)
   1306       return self._call_tf_sessionrun(
-> 1307           options, feed_dict, fetch_list, target_list, run_metadata)
   1308 

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\client\session.py in _call_tf_sessionrun(self, options, feed_dict, fetch_list, target_list, run_metadata)
   1408           self._session, options, feed_dict, fetch_list, target_list,
-> 1409           run_metadata)
   1410     else:

InvalidArgumentError: Matrix size-incompatible: In[0]: [1,16384], In[1]: [1024,10]
     [[Node: dense_251/MatMul = MatMul[T=DT_FLOAT, _class=["loc:@training_22/RMSprop/gradients/dense_251/MatMul_grad/MatMul"], transpose_a=false, transpose_b=false, _device="/job:localhost/replica:0/task:0/device:GPU:0"](flatten_153/Reshape, dense_251/kernel/read)]]
     [[Node: loss_26/mul/_579 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1108_loss_26/mul", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Can somebody help me figuring out what I am doing wrong here?

Answer 1

The exception message Keras/TensorFlow gives you is (to be honest) not as helpful as one might hope it to be.

One thing you always should check is: do I compute the output shape of my custom layer correctly? You are returning:

return (self.output_dim[0], self.output_dim[1], input_shape[2][3])

but this is totally ignoring that your data will be batched (as the shape only has rank 3). You can fix it by adding a None as the first item of the tuple:

return (None, self.output_dim[0], self.output_dim[1], input_shape[2][3])

While trying to find the real problem/solve your issue I noticed, that the code you referenced has some other problems, too. I fixed these as well; you can find a reimplemented version of the code in this repository.

PS: You could have noticed this problem already on your own, respectively you've already found a clue about it:

when I use the MaxPool2D that is commented out right now, than I get a "IndexError: tuple index out of range

this error message should have warned you, that the output shape of the layer might not be correct/as it is intended.