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.