如何在PyTorch中实现与tf.nn.conv2d_transpose效果一致的运算

核心差异说明 #

两种框架的转置卷积实现有两个核心区别需要适配：

• 数据排布差异：TensorFlow默认使用NHWC（批次、高度、宽度、通道）的4维张量排布，PyTorch默认使用NCHW（批次、通道、高度、宽度）排布
• 卷积核格式差异：tf.nn.conv2d_transpose的卷积核格式为[kernel_h, kernel_w, output_channels, input_channels]，PyTorchConvTranspose2d的卷积核格式为[input_channels, output_channels, kernel_h, kernel_w]

完整等效实现代码 #

import numpy as np
import tensorflow as tf
import torch

# 你给出的初始化代码
np.random.seed(42)
y_val = np.random.rand(1, 32, 32, 1024)
feats_val = np.random.rand(3, 3, 128, 1024)

y_tf = tf.Variable(y_val)
feats_tf = tf.Variable(feats_val)

y_tor = torch.tensor(y_val)
feats_tor = torch.tensor(feats_val)

# TensorFlow原有运算
y_up_tf = tf.nn.conv2d_transpose(y_tf, feats_tf, [1, 64, 64, 128], strides=[1,2,2,1])

# PyTorch等效运算
# 1. 调整输入张量从NHWC到NCHW格式
y_tor_nchw = y_tor.permute(0, 3, 1, 2).float()
# 2. 调整卷积核格式从TF规范转成PyTorch规范
feats_tor_perm = feats_tor.permute(3, 2, 0, 1).float()
# 3. 调用转置卷积，参数匹配TF配置
y_up_tor_nchw = torch.nn.functional.conv_transpose2d(
    y_tor_nchw,
    feats_tor_perm,
    stride=2,
    padding=1,
    output_padding=1
)
# 4. 输出转回NHWC格式和TF输出对齐
y_up_tor = y_up_tor_nchw.permute(0, 2, 3, 1)

# 验证数值一致性，输出为True即说明结果完全一致
print(np.allclose(y_up_tf.numpy(), y_up_tor.numpy(), atol=1e-6))

参数说明：这里的padding=1和output_padding=1是匹配示例中输出尺寸64x64的配置，符合转置卷积输出尺寸公式：output_size = input_size * stride + output_padding - 2*padding + kernel_size，代入数值刚好得到64的输出尺寸，和TensorFlow指定的输出维度完全匹配。

内容的提问来源于stack exchange，提问作者Ruisheng Lei