Hi, I’m attempting to port a model from PyTorch to TF 2.x/Keras. I was able to load PyTorch weights in Keras, but the outputs does not seem to be equal, even for a simple conv 2D layer.
The following is the code that I used.
def filter_dict_by_name(name, state_dict):
    return dict(filter(lambda t: t[0].startswith(name),  state_dict.items()))


def strip_dict_key_prefix(state_dict, cut=-2):
    """Strip parameter key prefix."""
    return {f"{'.'.join(k.split('.')[cut:])}": v for k, v in state_dict.items()}


def conv2d_weights_pth2tf(params_dict):
    """Convert conv2d weights from Pytorch to TensorFlow."""
    
    weights = []
    
    params_dict = strip_dict_key_prefix(params_dict, -1)
    
    # gamma, beta, moving_mean, moving_variance
    for key in ["weight", "bias"]:
        if key not in params_dict:
            continue
            
        if key == "weight":
            weights.append(np.transpose(params_dict[key].detach().numpy(), [2, 3, 1, 0]))
        else:
            weights.append(params_dict[key].detach().numpy())
    
    return weights


def get_weight_conversion_fn(layer):
    """Get the appropriate weight conversion function given the layer."""
    
    if isinstance(layer, keras.layers.Conv2D):
        fn = conv2d_weights_pth2tf
    elif isinstance(layer, keras.layers.BatchNormalization):
        fn = bn_weights_pth2tf
    else:
        raise Exception(f"Unknown layer type: {type(layer)}")
        
    return fn

# Create two conv 2D layers, one in TF, another in PyTorch
conv1_tf = keras.layers.Conv2D(
            64,
            kernel_size=7,
            strides=2,
            padding="same",
            use_bias=False
        )

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

conv1_state_dict = strip_dict_key_prefix(filter_dict_by_name("conv1", state_dict), -1)
conv1.load_state_dict(
    conv1_state_dict
)
# outputs: <All keys matched successfully>
conv1.eval()

conv1_tf.set_weights(
    get_weight_conversion_fn(conv1_tf)(
        conv1_state_dict
    )
)

conv1_out = conv1(torch.from_numpy(fixed_input).permute(0, 3, 1, 2)) \
                .permute(0, 2, 3, 1) \
                .detach() \
                .numpy()
conv1_out_tf = conv1_tf(fixed_input)\
                .numpy()

np.allclose(conv1_out, conv1_out_tf)
# outputs: False