Lambda layer runs twice even though called only once

[guyeng0]

I'm trying to run a tensorflow tf.nn.max_pool_with_argmax in order to use in an unpooling layer. I used a Lambda layer to run the tensorflow function. But it seems that even though I'm calling the Lambda layer only once, the function runs twice and outputs two argmax layers. Not sure if I'm doing something incorrectly or whether this is a bug. Also, I'm not an expert in python so probably the code could be written better. The inputs used come from google's notMNist 28x28 dataset.

import numpy as np
import tensorflow as tf
import os
from keras.layers import Convolution2D, Dense, Flatten, Input, Lambda
from keras.models import Model
from keras.optimizers import SGD

from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_nn_ops
@ops.RegisterGradient("MaxPoolWithArgmax")
def _MaxPoolWithArgmaxGrad(op, grad, some_other_arg):
return gen_nn_ops._max_pool_grad(op.inputs[0],
op.outputs[0],
grad,
op.get_attr("ksize"),
op.get_attr("strides"),
padding=op.get_attr("padding"),
data_format='NHWC')

def keras_max_pool_with_argmax(input):
global argmaxlayers
(maxpoolout, maxpoolargmax) = tf.nn.max_pool_with_argmax(input, [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME")
print ("adding")
argmaxlayers.append(maxpoolargmax)
return maxpoolout

argmaxlayers = []

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
inputs = Input(batch_shape=(100,28,28,1))
conv1 = Convolution2D(4, 3, 3, init='he_normal', activation = 'relu', border_mode='same', bias=True)(inputs)
maxpoolout1 = Lambda(keras_max_pool_with_argmax)(conv1)
flat = Flatten()(maxpoolout1)
predictions = Dense(10, init='he_normal', bias = True, activation = 'softmax', name='predictions')(flat)
model = Model(input=inputs, output=predictions)
sgd = SGD(lr=0.03, decay=0.0, momentum=0.0)
model.compile(loss='categorical_crossentropy', optimizer=sgd)
model.fit(train_dataset[0:100000], train_labels[0:100000], batch_size=100, nb_epoch=1, callbacks=[],
validation_data=(valid_dataset, valid_labels))
print (argmaxlayers)

In the output I see

adding
adding
Train on 100000 samples, validate on 10000 samples
Epoch 1/1
100000/100000 [==============================] - 2s - loss: 0.7490 - val_loss: 0.6428
[<tf.Tensor 'MaxPoolWithArgmax:1' shape=(100, 14, 14, 4) dtype=int64>, <tf.Tensor 'MaxPoolWithArgmax_1:1' shape=(100, 14, 14, 4) dtype=int64>]

So "adding" is outputted twice indicating the function keras_max_pool_with_argmax was called twice, and additionally I have two separate tensors that were created instead of just one.

Thanks,

Guy

[unrealwill]

Hello,
There are probably two graphs which are constructed. One for training. The other for validation.

You shouldn't use some global variable to save some internal state.

Create a custom layer which outputs two tensors. (They are created the same way that normal layers except you may have to additionally overload compute_mask (just make it return [None,None])

[guyeng0]

@unrealwill Thanks for your reply. I removed the validation data from model.fit and this still happens. Does the validation graph get created regardless of whether I feed it validation data? Are there two graphs created every time? This seems like a waste of memory since they should have identical weights in the layers.

In addition, even if two graphs are created, one for training and the other for validation, my plan is to use these argmax values in an unpooling layer later on. How do I know which one of the two tensors that were created to feed into the unpooling layer?

Regarding creating a custom layer to output two tensors. I didn't realize this was how I can do that. Thanks.

[unrealwill]

Weights are shared so this doesn't consume much more memory. I'm not sure whether or not they are created even if not used, but probably as it's cheap. It probably has to do with "LearningPhase".

[guyeng0]

And how do I know which of the tensors I should use for an unpooling layer? Is the first one in the list from the training graph and the second from the validation graph?

[unrealwill]

You don't have to worry about this. It's transparent. Do as if there is only one graph. But don't use global.

[guyeng0]

@unrealwill, thanks for all your help. I tried it and it works. I'll close the issue, but just in case someone else needs to I'm pasting my naively implemented layer which does tensorflow max_pool_with_argmax and outputs the two tensors.

class KerasMaxPooling2DWithArgmax(Layer):
def init(self, **kwargs):
super(KerasMaxPooling2DWithArgmax, self).init(**kwargs)

def call(self, x, mask=None):
    tensorlist = tf.nn.max_pool_with_argmax(x, (1, 2, 2, 1), strides = (1, 2, 2, 1), padding = "SAME")
    return [tensorlist[0], tensorlist[1]]

def get_output_shape_for(self, input_shape):
    return [(input_shape[0], input_shape[1]/2, input_shape[2]/2, input_shape[3]),(input_shape[0], input_shape[1]/2, input_shape[2]/2, input_shape[3])]

def compute_mask(self, input, input_mask=None):
    return [None, None]