I want to replace specific layers in pytorch for a given nn.Sequential.
Example: Replace maxpool with average pool
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.