How can I replace all layers with my own?

What is the correct method for replacing layers in the network with your own?

For example: I have a custom convolution layer and I want to change all conv layers in pretrained ResNet18.
Now I use the following code:

for name, p in net.named_parameters():
        if not ("layer" in name and ("conv" in name or "downsample" in name)):
            continue

        lnames = name.split('.')
        if not "module" in lnames[0]:
            if "downsample" in name:
                cnv = net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]]  # nn.Conv2d()
            else:
                cnv = net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]  # nn.Conv2d()
        else:
            if "downsample" in name:
                cnv = net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]]._modules[lnames[4]]  # nn.Conv2d()
            else:
                cnv = net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]] #nn.Conv2d()

        cnv_type = type(cnv).__name__
        if cnv_type != "Conv2d":
            continue

        custom_cnv = CustomConv2d(cnv.in_channels, cnv.out_channels, cnv.kernel_size,
								  best_th, best_scale, cnv.stride, cnv.padding, cnv.dilation,
								  cnv.groups, weights=cnv.weight.data)

        if not "module" in lnames[0]:
            if "downsample" in name:
                net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]]  = custom_cnv
            else:
                net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]] = custom_cnv
        else:
            if "downsample" in name:
                net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]]._modules[lnames[4]] = custom_cnv
            else:
                net._modules[lnames[0]]._modules[lnames[1]]._modules[lnames[2]]._modules[lnames[3]] = custom_cnv

But I think this method is bad.

Thank for your suggestion.

We could use add_module to replace your own custom function in PyTorch.

firstbn = resnet.layer1
firstbn[0].add_module('bn1', ABN(64))
firstbn[1].add_module('bn1', ABN(64))

ABN is my own class to replace bn1 normalization layer, if we use the same name bn1, it will replace the old module with our own custom module.