How to directly replace a certain layer in the pretrained ConvertModel in pytorch


I have a pytorch model that is converted from onnx model.

net = copy.deepcopy(model_ori)

And print(net) returns as follows:

  (Constant_9): Constant(constant=tensor([[[[0.13070001]]]], device='cuda:0'))
  (Sub_10): sub()
  (Constant_11): Constant(constant=tensor([[[[0.30810001]]]], device='cuda:0'))
  (Div_12): Div()
  (Conv_13): Conv2d(1, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  (Relu_14): ReLU(inplace=True)
  (Conv_15): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  (Relu_16): ReLU(inplace=True)
  (Flatten_17): Flatten()
  (Gemm_18): Linear(in_features=1568, out_features=1000, bias=True)
  (Relu_19): ReLU(inplace=True)
  (Gemm_output): Linear(in_features=1000, out_features=10, bias=True)

I wonder how I can directly replace one layer with another layer (eg: Conv_13 to another self-defined layer) and net is updated as well?
I appreciate any suggestions. Thank you!

Two considerations:

  1. Sizing issues - if you plan to replace the conv2d layer with one of a different size or layer type, that is probably not going to work, unless you make sure the input and output of this intermediate layer(s) are of both the same size and dimensions. In such a case, you’ll probably be better off defining the new structure with a new forward pass and then load your pretrained weights to each layer as appropriate, if transfer learning is involved.
  2. If you’re only changing the weights of Conv_13, you can just use with torch.no_grad(): and set the weights and/or bias equal to whatever pretrained tensor you have.

For example:

with torch.no_grad():

Thank you for your heads-up!
I used net._modules['Conv_13'] = new_layer to do the update

1 Like