tensorflow keras: use variable of one Layer in another (Serialization Error: TypeError: can't pickle _thread.RLock objects)

I want to use variables of a layer A in a subsequent layer B. This causes serialization issues depending on the scope.

• Concretely I have this general (minimal) setup:

X_in -> normalize(X_in) -> some_layers(X) -> unnormalize(X) -> X_out

• For the unnormalize(X) step, I need to use the mean and variance found in normalize(X)

• For this purpose, I am using tensorflow's Normalization Layer and its variables mean and variance

When I write everything in global scope, that setup works:

dat = np.random.random(size=(100,7))

########## 1: global scope ##########

input = tf.keras.Input(shape=dat.shape[1:])
norm_layer = tf.keras.layers.experimental.preprocessing.Normalization()
x = norm_layer(input)
output = tf.keras.layers.Lambda(lambda xx: xx * tf.sqrt(norm_layer.variance) + norm_layer.mean)(x)

model = tf.keras.Model(input,output)
model.compile(optimizer='adam',loss=tf.keras.losses.mse)
norm_layer.adapt(dat)
model.fit(dat,dat,batch_size=10,epochs=3,verbose=2)

model.save('test_model.h5')
print('\n', '*'*20, 'global scope: model saved', '*'*20)

but when the above code is in a function scope:

########## 2: function scope ##########

def function_scope(dat):
    input = tf.keras.Input(shape=dat.shape[1:])
    norm_layer = tf.keras.layers.experimental.preprocessing.Normalization()
    x = norm_layer(input)
    output = tf.keras.layers.Lambda(lambda xx: xx * tf.sqrt(norm_layer.variance) + norm_layer.mean)(x)

    model = tf.keras.Model(input,output)
    model.compile(optimizer='adam',loss=tf.keras.losses.mse)
    norm_layer.adapt(dat)
    model.fit(dat,dat,batch_size=10,epochs=3,verbose=2)

    model.save('test_model.h5')
    print('*'*20, 'function scope: model saved', '*'*20)


function_scope(dat)

I get the following error:

  File "/usr/lib/python3.7/copy.py", line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.7/copy.py", line 169, in deepcopy
    rv = reductor(4)
TypeError: can't pickle _thread.RLock objects

I checked the solutions here, but only one seems to work: capturing norm_layer.variance and norm_layer.mean in a function closure:

norm_layer = tf.keras.layers.experimental.preprocessing.Normalization()
x = norm_layer(input)
def unnnorm_closure(xx):
    return xx * tf.sqrt(norm_layer.variance) + norm_layer.mean
output = tf.keras.layers.Lambda(lambda xx: unnnorm_closure(xx))(x)

The lambda wrap doesn't work for me:

norm_layer = tf.keras.layers.experimental.preprocessing.Normalization()
x = norm_layer(input)
var, mean = tf.keras.layers.Lambda(lambda xx: (xx[0],xx[1]))((norm_layer.variance, norm_layer.mean))
output = tf.keras.layers.Lambda(lambda xx: xx * tf.sqrt(var) + mean)(x)

and neither does a custom layer:

class UnNormalization(tf.keras.layers.Layer):

    def __init__(self, norm_mean, norm_var, **kwargs):
        super(UnNormalization, self).__init__(**kwargs)
        self.norm_mean = norm_mean
        self.norm_var = norm_var
        self.unnormalize = tf.keras.layers.Lambda(lambda xx: xx * tf.sqrt(self.norm_var) + self.norm_mean)

    def call(self, inputs):
        return self.unnormalize(inputs)

    def get_config(self):
        config = super(UnNormalization, self).get_config()
        config.update({'norm_mean': self.norm_mean, 'norm_var': self.norm_var})
        return config

There is another related question here but it is unanswered. Although I managed to fix the issue with the closure, I would like to know what is happening. If I understand correctly, the main problem is that, at the point of serialization, norm_layer.mean and norm_layer.variance are Keras symbolic tensors.

I am using Python 3.7.3, Tensorflow 2.2.0