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