My code:
import tensorflow as tf
from tensorflow.keras.layers import Conv2D
import torch, torchvision
import torch.nn as nn
import numpy as np
# Define the PyTorch layer
pt_layer = torch.nn.Conv2d(3, 12, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
# Get the weight tensor from the PyTorch layer
pt_weights = pt_layer.weight.detach().numpy()
# Create the equivalent Keras layer
keras_layer = Conv2D(12, kernel_size=(3, 3), strides=(2, 2), padding='same', use_bias=False, input_shape=(None, None, 3))
# Build the Keras layer to initialize its weights
keras_layer.build((None, None, None, 3))
# Transpose the PyTorch weights to match the expected shape of the Keras layer
keras_weights = pt_weights.transpose(2, 3, 1, 0)
# Set the weights of the Keras layer to the PyTorch weights
keras_layer.set_weights([keras_weights])
#Test both models
arr = np.random.normal(0,1,(1, 3, 224, 224))
print(pt_layer(torch.from_numpy(arr).float())[0,0])
print(keras_layer(arr.transpose(0,2,3,1))[0,:,:,0])
I would expect both prints to be quite similar, but the are really different. I ran it on a Colab to make sure it wasn't due to old Pytorch/Keras versions. I'm sure I missed something trivial, bu I can't find it... Any help would be welcome, please.
Found the answer: The padding in Keras and Pytorch are quite different it seems. To fix, use ZeroPadding2D instead:
keras_layer = tf.keras.Sequential([
ZeroPadding2D(padding=(1, 1)),
Conv2D(12, kernel_size=(3, 3), strides=(2, 2), padding='valid', use_bias=False, input_shape=(None, None, 3))
])