TensorFlow条件GAN转PyTorch遇ConvTranspose2D维度错误求助

核心错误原因 #

1. 输入定义逻辑错误：PyTorch代码在__init__中将ops和inputs定义为nn.Parameter或空张量，这完全违背条件GAN的输入逻辑——这两个是模型的输入数据，不是可训练参数，仅需作为forward函数的传入参数，不需要在初始化阶段定义。
2. 缺少形状转换+通道维度位置不匹配：
   • TensorFlow中Dense层后通过Reshape((16,16,2))将扁平张量转为2D特征图，PyTorch代码注释掉了该步骤，导致conv_transpose_1接收到的仍是扁平的[batch_size, 512]张量，不符合ConvTranspose2D要求的4D输入格式[batch_size, channels, height, width]。
   • TensorFlow的通道维度在最后（(H,W,C)），而PyTorch要求通道维度在第二维（(C,H,W)），转换时必须调整维度顺序。
3. 卷积层参数与原代码不一致：原TensorFlow中所有Conv2DTranspose使用strides=1、padding="same"，但PyTorch代码里第一个conv_transpose_1写了stride=2，与原逻辑冲突。

修正后的完整代码 #

模型类实现 #

import torch
import torch.nn as nn

class CondGANGenerator(nn.Module):
    def __init__(self, h_size, num_points):
        super().__init__()
        self.initializer = nn.init.normal_

        # 定义网络层，不在这里定义输入数据
        self.x = nn.Sequential(
            nn.Linear(num_points, 16 * 16 * 2, bias=False),
            nn.LeakyReLU(),
            # 新增形状转换：先转成(batch, 16,16,2)，再调整通道到第二维
            nn.Unflatten(1, (16, 16, 2)),
            nn.Permute(0, 3, 1, 2)
        )

        # 对齐原TensorFlow的卷积参数：stride=1 + padding=1 等价于padding="same"
        self.conv_transpose_1 = nn.Sequential(
            nn.ConvTranspose2d(2, 64, kernel_size=4, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(64),
            nn.LeakyReLU(),
        )

        self.conv_transpose_2 = nn.Sequential(
            nn.ConvTranspose2d(64, 64, kernel_size=4, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(64),
            nn.LeakyReLU(),
        )

        self.conv_transpose_3 = nn.Sequential(
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=1, padding=1, bias=False),
        )

        self.conv_transpose_4 = nn.Sequential(
            nn.ConvTranspose2d(32, 2, kernel_size=4, stride=1, padding=1, bias=False),
        )

        # 替换为你实际的自定义层/函数
        self.density_matrix = lambda x: x  # 示例占位符
        self.expectation = lambda ops, x, pref: x  # 示例占位符
        self.noise = nn.GaussianNoise(0.1)  # 示例噪声层

    def forward(self, ops, inputs):
        # 输入形状要求：
        # ops: [batch_size, h_size, h_size, num_points*2]
        # inputs: [batch_size, num_points]
        x = self.x(inputs)
        # 此时x形状为[batch_size, 2, 16, 16]，符合卷积层输入要求
        x = self.conv_transpose_1(x)
        x = self.conv_transpose_2(x)
        x = self.conv_transpose_3(x)
        x = self.conv_transpose_4(x)
        x = self.density_matrix(x)
        complex_ops = convert_to_complex_ops(ops)
        prefactor = 1.0
        x = self.expectation(complex_ops, x, prefactor)
        x = self.noise(x)
        return x

关键修正点说明 #

• 输入逻辑修正：删除__init__中ops = nn.Parameter(...)和inputs = torch.empty(...)的错误定义，模型输入完全由forward参数传入。
• 形状转换修正：用nn.Unflatten将Linear层输出的扁平张量转为2D特征图，再用nn.Permute调整通道维度位置，得到PyTorch卷积层要求的格式。
• 卷积参数对齐：将第一个conv_transpose_1的stride=2改为stride=1，确保和原TensorFlow代码的输出形状、逻辑一致。

调用示例 #

# 初始化模型（对应原TensorFlow的h_size=16，num_points=1296）
model = CondGANGenerator(h_size=16, num_points=1296)
# 生成测试输入
batch_size = 8
ops = torch.randn(batch_size, 16, 16, 1296*2)
inputs = torch.randn(batch_size, 1296)
# 前向传播测试
output = model(ops, inputs)
print(output.shape)

内容的提问来源于stack exchange，提问作者Silvio sjsj