<jupyter_start><jupyter_text>Copyright 2019 Google LLCLicensed under the Apache License, Version 2.0 (the "License");you may not use this file except in compliance with the License.You may obtain a copy of the License at    https://www.apache.org/licenses/LICENSE-2.0Unless required by applicable law or agreed to in writing, softwaredistributed under the License is distributed on an "AS IS" BASIS,WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.See the License for the specific language governing permissions andlimitations under the License.  Example of folding keras conv layer with batch norm  Imports<jupyter_code>import tensorflow as tf
import numpy as np
np.random.seed(123)<jupyter_output><empty_output><jupyter_text>Model with convolution and batch norm layer definition<jupyter_code>epsilon=0.001
inputs = tf.keras.Input(shape=(50, 32, 5), batch_size=4)
net = inputs
net = tf.keras.layers.Conv2D(filters=2, kernel_size=(3,3))(net)
net = tf.keras.layers.BatchNormalization(epsilon=epsilon)(net)
net = tf.keras.layers.ReLU()(net)
net = tf.keras.layers.Flatten()(net)
model = tf.keras.Model(inputs, net)
model.summary()
model.layers<jupyter_output><empty_output><jupyter_text>Initialize all model weights with random numbers<jupyter_code># we will set all weights to randmom numbers so that even bias will be non zero 
# it will help to validated numerical correctness of conv and batch norm fusion
all_weights = model.get_weights()
for i in range(len(all_weights)):
  all_weights[i] = np.random.random(all_weights[i].shape)
model.set_weights(all_weights)<jupyter_output><empty_output><jupyter_text>Dims of conv layer weights<jupyter_code>ind_conv_layer = 1
assert(isinstance(model.layers[ind_conv_layer], tf.keras.layers.Conv2D))
conv_weights = model.layers[ind_conv_layer].get_weights()
print("conv weights shape " + str(conv_weights[0].shape))
print("conv bias shape " + str(conv_weights[1].shape))<jupyter_output>conv weights shape (3, 3, 5, 2)
conv bias shape (2,)<jupyter_text>Dims of batch norm layer weights<jupyter_code>ind_batch_norm_layer = 2
assert(isinstance(model.layers[ind_batch_norm_layer], tf.keras.layers.BatchNormalization))
bn_weights = model.layers[ind_batch_norm_layer].get_weights()
gamma = bn_weights[0]
print("gamma shape " + str(gamma.shape))
betta = bn_weights[1]
print("betta shape " + str(gamma.shape))
mean = bn_weights[2]
print("mean shape " + str(gamma.shape))
variance = bn_weights[3]
print("variance shape " + str(gamma.shape))<jupyter_output>gamma shape (2,)
betta shape (2,)
mean shape (2,)
variance shape (2,)<jupyter_text>Fuse conv and batch norm weights<jupyter_code>new_conv_weights = np.multiply(conv_weights[0], gamma) / np.sqrt(variance + epsilon)
new_bias = betta + np.multiply((conv_weights[1] - mean), gamma) / np.sqrt(variance + epsilon)<jupyter_output><empty_output><jupyter_text>Model with folded/fused convolution and batch norm layers<jupyter_code>inputs = tf.keras.Input(shape=(50, 32, 5), batch_size=4)
net = inputs
net = tf.keras.layers.Conv2D(filters=2, kernel_size=(3,3))(net)
net = tf.keras.layers.Flatten()(net)
model_fused = tf.keras.Model(inputs, net)
model_fused.summary()<jupyter_output>Model: "functional_25"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_17 (InputLayer)        [(4, 50, 32, 5)]          0         
_________________________________________________________________
conv2d_15 (Conv2D)           (4, 48, 30, 2)            92        
_________________________________________________________________
flatten_13 (Flatten)         (4, 2880)                 0         
=================================================================
Total params: 92
Trainable params: 92
Non-trainable params: 0
_________________________________________________________________<jupyter_text>Initialize model_fused with fused weights<jupyter_code>all_weights_fused = model_fused.get_weights()
all_weights_fused[0] = new_conv_weights
all_weights_fused[1] = new_bias
model_fused.set_weights(all_weights_fused)<jupyter_output><empty_output><jupyter_text>Validate that model and model_fused produce the same outputs<jupyter_code>input_data = np.random.random(inputs.shape)
outputs = model.predict(input_data)
outputs_fused = model_fused.predict(input_data)
np.allclose(outputs, outputs_fused)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1278549037169727, 'fraction_numerical': 0.020824003582624272, 'mean_word_length': 6.49496644295302, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 34, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '392503', 'license': 'Apache-2.0', 'n_tokens_mistral': 1444, 'n_tokens_neox': 1369, 'n_words': 330, 'path': 'kws_streaming/colab/conv_batch_norm_folding.ipynb', 'repo_name': 'ojInc/google-research', 'size': '4466'}