Tensorflow op as Lambda layer breaks serialization

I’m trying to use a tf op inside a Lambda layer. The goal is to basically wrap tf layers such that model.compile etc. still work. As I understood it I can just incorporate tf ops inside Lambda layers. If I do this however, serialization seems to be broken because it fails to pickle:

TypeError: can't pickle _thread.lock objects

The weird thing is that if I create a model in the main method, it seems to work fine. If I create the model inside a separate method, it fails to serialize. I’m guessing this is due to some scoping problem, but it could also be a Keras issue. A minimum example where this occurs is shown below, where I replaced a Conv2D with a tf conv2d op to demonstrate the issue.

import tensorflow as tf

from keras.models import Model
from keras.layers import Input, Lambda

def weight_variable(shape):
	initial = tf.truncated_normal(shape, stddev=0.1)
	return tf.Variable(initial)

def conv2d(x, W):
	return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def create_model(num_classes):
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)
	model = Model(inputs=inputs, outputs=x)
	return model

if __name__=='__main__':
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)

	# This works
	model_working = Model(inputs=inputs, outputs=x)
	print(model_working.to_json())

	# This gives "TypeError: can't pickle _thread.lock objects"
	model_broken = create_model(10)
	print(model_broken.to_json())

edit: to_json, to_yaml and save are at least broken with this code. I’m guessing its all methods that require the serialization of the network architecture that are broken with this model.

[jlin816]

@hgaiser, did you ever find a workaround for this issue? I’m encountering the same problem, even in the main method.

[hgaiser]

@minda163 https://github.com/fizyr/keras-maskrcnn/blob/master/keras_maskrcnn/layers/misc.py#L22-L27