Replace specific layer

I want to replace specific layers in pytorch for a given nn.Sequential.

Example: Replace maxpool with average pool

>>> model = nn.Sequential(nn.Conv2d(3,14,3),
                 nn.MaxPool2d((2,2)),
                 nn.Conv2d(14,28,3))
>>> print(model)
Sequential(
  (0): Conv2d(3, 14, kernel_size=(3, 3), stride=(1, 1))
  (1): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)
  (2): Conv1d(1, 2, kernel_size=(4,), stride=(1,))
)

The MaxPool layer is numbered (1).

>>> print(m[1])
MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)

Now change it by simple assignment.

>>> m[1] = nn.AvgPool2d((2,2))

Check the model again.

>>> print(m)
Sequential(
  (0): Conv2d(3, 14, kernel_size=(3, 3), stride=(1, 1))
  (1): AvgPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0)
  (2): Conv1d(1, 2, kernel_size=(4,), stride=(1,))
)

Hope it helps.