[FIXED] Manualy convert pytorch weights to tf.keras weights for convolutional layer

Issue

I'm trying to convert pytorch model to tf.keras model including weights conversion and came across an output missmatch between libraries' outputs.

Here I define two convolutional layers, which should be identical

torch_layer = torch.nn.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=(7, 7),
    stride=(2, 2),
    padding=(3, 3),
    dilation=1,
    groups=1,
    bias=False,
    padding_mode='zeros'
)

tf_layer = tf.keras.layers.Conv2D(
    filters=64,
    kernel_size=(7, 7),
    strides=(2, 2),
    padding='same',
    dilation_rate=(1, 1),
    groups=1,
    activation=None,
    use_bias=False
)
# define model to specify input channel size
tf_model = tf.keras.Sequential([tf.keras.layers.Input((256, 256, 3), batch_size=1), tf_layer])

now I have torch weights and I convert them to tf.keras format

# output_channels, input_channels, x, y
torch_weights = np.random.rand(64, 3, 7, 7)
# x, y, input_channels, output_channels
tf_weights = np.transpose(torch_weights, (2, 3, 1, 0))

# assign weights
torch_layer.weight = torch.nn.Parameter(torch.Tensor(torch_weights))
tf_model.layers[0].set_weights([tf_weights])

now I define input and the outputs are different (shape is the same, values are different), what am I doing wrong?

torch_inputs = np.random.rand(1, 3, 256, 256)
tf_inputs = np.transpose(torch_inputs, (0, 2, 3, 1))

torch_output = torch_layer(torch.Tensor(torch_inputs))
tf_output = tf_model.layers[0](tf_inputs)

Solution

In tensorflow, set_weights is basically used for outputs from get_weights, so it is better to use assign to avoid making mistakes.

Besides, 'same' padding in tensorflow is a little bit complicated. For details, see my SO answer. It depends on input_shape, kernel_size and strides. In your example here, it is translated to torch.nn.ZeroPad2d((2,3,2,3)) in pytorch.

Example codes: from tensorflow to pytorch

np.random.seed(88883)

#initialize the layers respectively
torch_layer = torch.nn.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=(7, 7),
    stride=(2, 2),
    bias=False
)
torch_model = torch.nn.Sequential(
              torch.nn.ZeroPad2d((2,3,2,3)),
              torch_layer
              )

tf_layer = tf.keras.layers.Conv2D(
    filters=64,
    kernel_size=(7, 7),
    strides=(2, 2),
    padding='same',
    use_bias=False
)

#setting weights in torch layer and tf layer respectively
torch_weights = np.random.rand(64, 3, 7, 7)
tf_weights = np.transpose(torch_weights, (2, 3, 1, 0))

with torch.no_grad():
  torch_layer.weight = torch.nn.Parameter(torch.Tensor(torch_weights))

tf_layer(np.zeros((1,256,256,3)))
tf_layer.kernel.assign(tf_weights)

#prepare inputs and do inference
torch_inputs = torch.Tensor(np.random.rand(1, 3, 256, 256))
tf_inputs = np.transpose(torch_inputs.numpy(), (0, 2, 3, 1))

with torch.no_grad():
  torch_output = torch_model(torch_inputs)
tf_output = tf_layer(tf_inputs)

np.allclose(tf_output.numpy() ,np.transpose(torch_output.numpy(),(0, 2, 3, 1))) #True

Edit: from pytorch to tensorflow

torch_layer = torch.nn.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=(7, 7),
    stride=(2, 2),
    padding=(3, 3),
    bias=False
)

tf_layer=tf.keras.layers.Conv2D(
    filters=64,
    kernel_size=(7, 7),
    strides=(2, 2),
    padding='valid',
    use_bias=False
    )

tf_model = tf.keras.Sequential([
           tf.keras.layers.ZeroPadding2D((3, 3)),
           tf_layer
           ])

Answered By - Laplace Ricky