如何用PyTorch构建音频条件GAN？生成样本白噪声问题排查

我正在构建一个可生成0-9数字人声音频的条件GAN，训练数据是已转换为Numpy数组的WAV文件，用librosa可正常回放清晰的数字发音。模型基于MNIST条件GAN修改，仅调整了输入形状，运行无报错，但训练时生成的样本回放为白噪声。我需要确定是模型未学习到有效特征，还是生成过程中样本被损坏。现有2700个尺寸为(256,64)的样本，相关代码如下：

数据加载代码 #

X_train = torch.tensor(X_train)
y_train = torch.tensor(y_train)
print("X_train Shape:", X_train.shape)
print("y_train Shape:", y_train.shape)
dataset = torch.utils.data.TensorDataset(X_train, y_train)

# 输出结果：
# X_train Shape: torch.Size([2700, 1, 256, 64])
# y_train Shape: torch.Size([2700])

loader = DataLoader(
    dataset,
    batch_size=30,
    shuffle=False,
)

判别器定义 #

class Discriminator(nn.Module):
    def __init__(self, channels, features_d, num_classes, x1_size, x2_size):
        super(Discriminator, self).__init__()
        self.x1_size = x1_size
        self.x2_size= x2_size
        self.disc = nn.Sequential(
            # input: N x channels x 64 x 64
            nn.Conv2d(channels+1, features_d, kernel_size=4, stride=2, padding=1), #in_channels = channels+1
            nn.LeakyReLU(0.2),
            # _block(in_channels, out_channels, kernel_size, stride, padding)
            self._block(features_d, features_d * 2, 4, 2, 1),
            self._block(features_d * 2, features_d * 4, 4, 2, 1),
            self._block(features_d * 4, features_d * 8, 4, 2, 1),
            # After all _block output is 4x4 (Conv2d below makes into 1x1)
            nn.Conv2d(features_d * 8, 1, kernel_size=4, stride=2, padding=0),
        )
        self.embed = nn.Embedding(num_classes, x1_size*x2_size) # Adds an additional channel

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False,
            ),
            nn.InstanceNorm2d(out_channels, affine=True),
            nn.LeakyReLU(0.2),
        )

    def forward(self, x, labels):
        embedding = self.embed(labels).view(labels.shape[0], 1, self.x1_size, self.x2_size) # for sound, size will need to be 2 variables
        x = torch.cat([x, embedding], dim=1) # Num_samples x c x size(h) x size(w)
        return self.disc(x)

生成器定义 #

class Generator(nn.Module):
    def __init__(self, channels_noise, channels_img, features_g, num_classes, x1_size, x2_size, embed_size):
        super(Generator, self).__init__()
        self.x1_size = x1_size
        self.x1_size = x1_size
        self.net = nn.Sequential(
            # Input: N x channels_noise x 1 x 1
            self._block(channels_noise+embed_size, features_g * 16, (16,4), 1, 0),  # img: 4x1
            self._block(features_g * 16, features_g * 8, 4, 2, 1),  # img: 16x4
            self._block(features_g * 8, features_g * 4, 4, 2, 1),  # img: 32x8
            self._block(features_g * 4, features_g * 2, 4, 2, 1),  # img: 64x16
            nn.ConvTranspose2d(
                features_g * 2, channels_img, kernel_size=4, stride=2, padding=1
                
            ),
             #Output: N x channels_img x 64 x 64
            nn.Tanh(),
        )
        self.embed = nn.Embedding(num_classes, embed_size)
    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False,
            ),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )

    def forward(self, x, labels):
        # latent vector z = num_samples x noise_dim x 1 x 1
        embedding = self.embed(labels).unsqueeze(2).unsqueeze(3)
        x = torch.cat([x, embedding], dim=1) 
        return self.net(x).to(device)

辅助函数 #

def gradient_penalty(critic, labels, real, fake, device="cpu"):
    BATCH_SIZE, C, H, W = real.shape
    alpha = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device)
    interpolated_images = real * alpha + fake * (1 - alpha)

    # Calculate critic scores
    mixed_scores = critic(interpolated_images, labels)

    # Take the gradient of the scores with respect to the images
    gradient = torch.autograd.grad(
        inputs=interpolated_images,
        outputs=mixed_scores,
        grad_outputs=torch.ones_like(mixed_scores),
        create_graph=True,
        retain_graph=True,
    )[0]
    gradient = gradient.view(gradient.shape[0], -1)
    gradient_norm = gradient.norm(2, dim=1)
    gradient_penalty = torch.mean((gradient_norm - 1) ** 2)
    return gradient_penalty


def save_checkpoint(state, filename="celeba_wgan_gp.pth.tar"):
    print("=> Saving checkpoint")
    torch.save(state, filename)


def load_checkpoint(checkpoint, gen, disc):
    print("=> Loading checkpoint")
    gen.load_state_dict(checkpoint['gen'])
    disc.load_state_dict(checkpoint['disc'])

训练循环代码 #

gen.train()
critic.train()

for epoch in range(NUM_EPOCHS):
    
    for batch_idx, (real, labels) in enumerate(loader):
        #print(batch_idx)
        real = real.to(device)
        cur_batch_size = real.shape[0]
        labels = labels.to(device)
        np.save(f"/content/drive/My Drive/Pytorch_Real/Real_array_{epoch}", real.detach().cpu().numpy())
        # Train Critic: max E[critic(real)] - E[critic(fake)]
        # equivalent to minimizing the negative of that
        for _ in range(CRITIC_ITERATIONS):
            noise = torch.randn(cur_batch_size, Z_DIM, 1, 1).to(device)
            fake = gen(noise, labels).to(device)
            np.save(f"/content/drive/My Drive/Pytorch_Fake/Fake_array_{epoch}", fake.detach().cpu().numpy())
            critic_real = critic(real, labels).reshape(-1).to(device)
            critic_fake = critic(fake, labels).reshape(-1).to(device)
            gp = gradient_penalty(critic, labels, real, fake, device=device)
            loss_critic = (
                -(torch.mean(critic_real) - torch.mean(critic_fake)) + LAMBDA_GP * gp
            )
            critic.zero_grad()
            loss_critic.backward(retain_graph=True)
            opt_critic.step()

        # Train Generator: max E[critic(gen_fake)] <-> min -E[critic(gen_fake)]
        gen_fake = critic(fake, labels).reshape(-1)
        loss_gen = -torch.mean(gen_fake)
        gen.zero_grad()
        loss_gen.backward()
        opt_gen.step()

        # Print losses occasionally and print to tensorboard
        if batch_idx % 10 == 0 and batch_idx > 0:
            print(
                f"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(loader)} \n                  Loss D: {loss_critic:.4f}, loss G: {loss_gen:.4f}"
            )

            with torch.no_grad():
                fake = gen(noise, labels)
                # take out (up to) 32 examples
                img_grid_real = torchvision.utils.make_grid(real[:32], normalize=True)
                img_grid_fake = torchvision.utils.make_grid(fake[:32], normalize=True)

                writer_real.add_image("Real", img_grid_real, global_step=step)
                writer_fake.add_image("Fake", img_grid_fake, global_step=step)

            step += 1

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