手動將 pytorch 權重轉換為卷積層的 tf.keras 權重

Question

我正在嘗試將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)

Answer 1

在 tensorflow 中， set_weights基本上是用於get_weights的輸出，所以最好使用assign以避免出錯。

此外，張量流中的“相同”填充有點復雜。 有關詳細信息，請參閱我的SO 答案。 它取決於input_shape 、 kernel_size和strides 。 在您的示例中，它被轉換為 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
           ])