[Python] Keras Functional replace intermediate layers

I would like to replace BatchNorm layers with GroupNorm in built-in keras models, e.g. ResNet50. I'm trying to reset nodes' layers to my new layer, however nothing changes when I query a model.summary().

import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras import layers

model = tf.keras.applications.resnet.ResNet50(include_top=False, weights=None)
channels = 3

for i,layer in enumerate(model.layers[:]):
if 'bn' in layer.name:
inbound_nodes = layer.inbound_nodes
outbound_nodes = layer.outbound_nodes

new_name = layer.name.replace('bn','gn')
new_layer = tfa.layers.GroupNormalization(channels)
new_layer._name = new_name

for j in range(len(inbound_nodes)):
inbound_nodes[j].layer = new_layer #set end of node to this layer

for k in range(len(outbound_nodes)):
new_layer.outbound_nodes.append(outbound_nodes[k])

layer = new_layer

Continue reading...