How to Replace a layer or Module in a pretrained network?

16

Just in case people find this useful, you can replace specific layers in a pretrained network with your customed layer iteratively as follow (or modify it according to your need).

def replace_module(module, target_name, new_module):
    for child_name, child_module in module.named_children():
        if target_name in child_name:
            setattr(module, child_name, new_module)
        else:
            replace_module(child_module, target_name, new_module)

# example:
t5_model = T5ForConditionalGeneration.from_pretrained(...)

custome_t5layernorm = CustomeT5LayerNorm(...)

replace_module(t5_model, 'layer_norm', custome_t5layernorm)