PyTorch nn.Conv2D to Flax linen.Conv

[kamalkraj]

Hi,

I am trying to move a model from PyTorch to Flax.

Is this correct translation for the 2D convolution layer from PyTorch to Flax

PyTorch

import torch.nn as nn
conv = nn.Conv2d(1, 32, kernel_size=(41, 11), stride=(2, 2), padding=(20, 5))

Flax

from flax import linen as nn
conv = nn.Conv(features=32, kernel_size=(41, 11), strides=(2, 2), padding=((0, 20), (0, 5)))

I am confused about the padding parameter?

[marcvanzee]

Have your checked the documentation? I think it is actually quite clear:

• From the Pytorch documentation: "padding controls the amount of padding applied to the input. It can be either a string {‘valid’, ‘same’} or a tuple of ints giving the amount of implicit padding applied on both sides."
• From the Flax documentation: "either the string ‘SAME’, the string ‘VALID’, the string ‘CIRCULAR’` (periodic boundary conditions), or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension."

So the torch Conv2d you provided above can be translated to Flax as follows:

 nn.Conv(features=32, kernel_size=(41, 11), strides=(2, 2), padding=((20, 20), (5, 5)))

[andsteing]

If you're porting PyTorch code you might want to check out this thread:
#1770 (reply in thread)