手动将 pytorch 权重转换为卷积层的 tf.keras 权重

时间:2021-06-28 14:42:35

标签: python tensorflow keras pytorch

我正在尝试将 pytorch 模型转换为 tf.keras 模型,包括权重转换,但遇到图书馆输出之间的输出不匹配。

这里我定义了两个卷积层,应该是一样的

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])

现在我有了火炬重量并将它们转换为 tf.keras 格式

# 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])

现在我定义输入和输出不同(形状相同,值不同),我做错了什么?

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)

1 个答案:

答案 0 :(得分:1)

在tensorflow中,set_weights基本用于get_weights的输出,所以最好使用assign避免出错。

此外,tensorflow 中的“相同”填充有点复杂。有关详细信息,请参阅我的 SO answer。这取决于 input_shapekernel_sizestrides。在您的示例中,它在 pytorch 中被转换为 torch.nn.ZeroPad2d((2,3,2,3))

示例代码:从 tensorflow 到 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

编辑:从pytorch到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
           ])