自定义文本嵌入条件GAN报错：running_mean应含1元素而非200

问题描述 #

实现基于文本嵌入的条件GAN时，Generator的embedding_layers中的BatchNorm1d层触发如下错误：

RuntimeError: running_mean should contain 1 elements not 200

数据集的图像形状为torch.Size([3, 512, 512])，文本嵌入形状为torch.Size([1, 768])，使用的超参数：embedding_dim=768、img_size=512、latent_dim=200、batch_size=32。

代码实现 #

Generator类 #

import torch.nn as nn
class Generator(nn.Module):
    def __init__(self, embedding_dim=300, latent_dim=100, image_size=64, num_channels=3):
        super(Generator, self).__init__()

        self.embedding_size = embedding_dim
        self.latent_dim = latent_dim
        self.image_size = image_size

        # Define embedding processing layers
        self.embedding_layers = nn.Sequential(
            nn.Linear(embedding_dim,latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.LeakyReLU(0.2, inplace=True)
        )

        # Define noise processing layers
        self.noise_layers = nn.Sequential(
            nn.Linear(latent_dim, image_size * image_size * 4),
            nn.BatchNorm1d(image_size * image_size * 4),
            nn.LeakyReLU(0.2, inplace=True)
        )

        # Define image processing layers
        self.conv_layers = nn.Sequential(
            nn.ConvTranspose2d(latent_dim + 256, 512, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, num_channels, kernel_size=4, stride=2, padding=1, bias=False),
            nn.Tanh()
        )
    def get_latent_dim(self):
        return self.latent_dim
    def forward(self, embeddings,noise):
        # Process embedding
        embedding_features = self.embedding_layers(embeddings)

        # Process noise
        
        noise_features = self.noise_layers(noise)

        # Combine features
        features = torch.cat((embedding_features, noise_features), dim=1)
        features = features.view(features.shape[0], -1, self.image_size // 16, self.image_size // 16)

        # Generate image
        image = self.conv_layers(features)

        return image

Discriminator类 #

import torch.nn as nn
class Discriminator(nn.Module):
    def __init__(self, embedding_dim=300, image_size=64, num_channels=3):
        super(Discriminator, self).__init__()

        # Define image processing layers
        self.conv_layers = nn.Sequential(
            nn.Conv2d(num_channels, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1024, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(1024, 1, kernel_size=4, stride=1, padding=0, bias=False),
            nn.Sigmoid()
        )

        # Define embedding processing layers
        self.embedding_layers = nn.Sequential(
            nn.Linear(embedding_dim, image_size * image_size),
            nn.BatchNorm1d(image_size * image_size),
            nn.LeakyReLU(0.2, inplace=True)
        )

    def forward(self, images, embeddings):
        # Process image
        image_features = self.conv_layers(images)

        # Process embedding
        embedding_features = self.embedding_layers(embeddings)
        embedding_features = embedding_features.view(embedding_features.shape[0], 1, 64, 64)

        # Combine features
        features = torch.cat((image_features, embedding_features), dim=1)

        # Classify
        classification = self.classification_layers(features).view(features.shape[0], -1)
        validity = self.validity_layers(features).view(features.shape[0], -1)

        return validity, classification

训练函数 #

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from tqdm import tqdm

def train_gan(generator, discriminator, dataset, batch_size, num_epochs, device):
    """
    Trains a conditional GAN with a generator and a discriminator using a PyTorch dataset containing text embeddings and images.
    """
    # Set up loss functions and optimizers
    adversarial_loss = nn.BCELoss()
    generator_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
    discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
    
    # Set up data loader
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    generator.to(device)
    discriminator.to(device)
    # Train the GAN
    for epoch in range(num_epochs):
        for i, data in enumerate(tqdm(data_loader)):
            # Load data and labels onto the device
            text_embeddings = data['text_embedding'].to(device)
            real_images = data['image'].to(device)
            
            # Generate fake images using the generator and the text embeddings
            noise = torch.randn(batch_size,generator.latent_dim).to(device)
            fake_images = generator(text_embeddings,noise)
            
            # Train the discriminator
            discriminator_optimizer.zero_grad()
            real_labels = torch.ones(real_images.size(0), 1).to(device)
            fake_labels = torch.zeros(fake_images.size(0), 1).to(device)
            real_predictions = discriminator(real_images, text_embeddings)
            real_loss = adversarial_loss(real_predictions, real_labels)
            fake_predictions = discriminator(fake_images.detach(), text_embeddings)
            fake_loss = adversarial_loss(fake_predictions, fake_labels)
            discriminator_loss = real_loss + fake_loss
            discriminator_loss.backward()
            discriminator_optimizer.step()
            
            # Train the generator
            generator_optimizer.zero_grad()
            fake_predictions = discriminator(fake_images, text_embeddings)
            generator_loss = adversarial_loss(fake_predictions, real_labels)
            generator_loss.backward()
            generator_optimizer.step()
            
            # Save generated images and model checkpoints every 500 batches
            if i % 500 == 0:
                with torch.no_grad():
                    fake_images = generator(text_embeddings[:16]).detach().cpu()
                save_image(fake_images, f"images\\generated_images_epoch_{epoch}_batch_{i}.png", normalize=True, nrow=4)
                torch.save(generator.state_dict(), f"images\\generator_checkpoint_epoch_{epoch}_batch_{i}.pt")
                torch.save(discriminator.state_dict(), f"images\\discriminator_checkpoint_epoch_{epoch}_batch_{i}.pt")
                
        # Print loss at the end of each epoch
        print(f"Epoch [{epoch+1}/{num_epochs}] Discriminator Loss: {discriminator_loss.item()}, Generator Loss: {generator_loss.item()}")

主函数 #

# defining hyperparamter 
torch.cuda.empty_cache()
embedding_dim=768
img_size=512
latent_dim=200
batch_size=32
num_epochs=100
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#define the main components
generator=Generator(embedding_dim=embedding_dim, latent_dim=latent_dim, image_size=img_size)
discriminator=Discriminator(embedding_dim=embedding_dim,image_size=img_size)

train_gan(generator=generator,
          discriminator=discriminator,
          dataset=dataset,
          batch_size=batch_size,
          num_epochs=num_epochs,
          device=device,)

错误原因分析 #

这个错误的核心是输入BatchNorm1d的张量维度不匹配：

• BatchNorm1d要求输入张量形状为[batch_size, num_features]，会对num_features维度做归一化。
• 你的文本嵌入原始形状是[1,768]，经过DataLoader加载后，形状变为[32,1,768]（多了一个冗余维度），而非预期的[32,768]。
• 这个冗余维度导致nn.Linear(768,200)输出形状为[32,1,200]，输入BatchNorm1d(200)时，层会误将第二维（尺寸为1）当作特征维度，与初始化的200不匹配，触发错误。

解决方法 #

1. 修正文本嵌入的维度 #

在训练函数中，挤压文本嵌入的冗余维度，确保输入到Generator的张量形状为[batch_size, embedding_dim]：

# 替换原代码中的文本嵌入加载行
text_embeddings = data['text_embedding'].squeeze(1).to(device)

2. 修复Generator的卷积层输入通道不匹配 #

原Generator的conv_layers第一个ConvTranspose2d输入通道设为latent_dim + 256，但实际拼接后的特征维度是latent_dim + image_size*image_size*4（当image_size=512时，这个值远大于前者），需修改输入通道：

# 替换原conv_layers的第一个层
nn.ConvTranspose2d(latent_dim + image_size*image_size*4, 512, kernel_size=4, stride=2, padding=1, bias=False),

更合理的优化：调整noise_layers的输出维度，让拼接后的特征维度更紧凑，避免过大的通道数导致显存压力。

3. 修复Discriminator的未定义层错误 #

原Discriminator的forward方法调用了未定义的self.classification_layers和self.validity_layers，需在__init__中补充定义：

# 在Discriminator的__init__中添加
self.classification_layers = nn.Sequential(
    nn.Conv2d(2, 64, kernel_size=3, stride=1, padding=1),
    nn.LeakyReLU(0.2),
    nn.Flatten(),
    nn.Linear(64*13*13, 10)  # 13是conv_layers输出的特征图尺寸，根据实际计算调整
)
self.validity_layers = nn.Sequential(
    nn.Conv2d(2, 64, kernel_size=3, stride=1, padding=1),
    nn.LeakyReLU(0.2),
    nn.Flatten(),
    nn.Linear(64*13*13, 1)
)

同时，需调整embedding_features的尺寸，使其与image_features的特征图尺寸匹配：

# 替换Discriminator forward中的embedding_features.view行
embedding_features = embedding_features.view(embedding_features.shape[0], 1, 13, 13)

4. 修复训练函数中的Generator调用错误 #

保存图像时，Generator的forward需要两个参数（embeddings和noise），需补充噪声输入：

# 替换训练函数中保存图像的代码块
if i % 500 == 0:
    with torch.no_grad():
        noise = torch.randn(16, generator.latent_dim).to(device)
        fake_images = generator(text_embeddings[:16], noise).detach().cpu()
    save_image(fake_images, f"images\\generated_images_epoch_{epoch}_batch_{i}.png", normalize=True, nrow=4)
    torch.save(generator.state_dict(), f"images\\generator_checkpoint_epoch_{epoch}_batch_{i}.pt")
    torch.save(discriminator.state_dict(), f"images\\discriminator_checkpoint_epoch_{epoch}_batch_{i}.pt")

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