训练GAN模型时遭遇TensorFlow remapper失败错误求助

错误信息 #

2024-02-06 15:10:19.270018: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:961] remapper failed: INVALID_ARGUMENT: Mutation::Apply error: fanout 'gradient_tape/sequential/leaky_re_lu_1/LeakyRelu_1/LeakyReluGrad' exist for missing node 'sequential/conv2d/BiasAdd_1'.

问题场景 #

训练自定义GAN模型时触发该错误，原模型代码如下：

def build_generator(): 
    model = Sequential()

    # Takes in random values
    # Beginnings of a generated image
    model.add(Dense(16*16*512, input_dim=512))
    model.add(LeakyReLU(0.2))
    model.add(Reshape((16,16,512)))

    # Upsampling block 1 
    model.add(UpSampling2D())
    model.add(Conv2D(512, 5, padding='same'))
    model.add(LeakyReLU(0.2))

    # Upsampling block 2 
    model.add(UpSampling2D())
    model.add(Conv2D(512, 5, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(UpSampling2D())
    model.add(Conv2D(512, 5, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(UpSampling2D())
    model.add(Conv2D(512, 5, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(UpSampling2D())
    model.add(Conv2D(512, 5, padding='same'))
    model.add(LeakyReLU(0.2))    
    # Convolutional block 1
    model.add(Conv2D(512, 4, padding='same'))
    model.add(LeakyReLU(0.2))

    # Convolutional block 2
    model.add(Conv2D(512, 4, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(Conv2D(512, 4, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(Conv2D(512, 4, padding='same'))
    model.add(LeakyReLU(0.2))

    model.add(Conv2D(512, 4, padding='same'))
    model.add(LeakyReLU(0.2))

    # Conv layer to get to one channel
    model.add(Conv2D(3, 4, padding='same', activation='sigmoid'))

    return model

def build_discriminator(): 
    model = Sequential()

    # First Conv Block
    model.add(Conv2D(32, 5, input_shape = (512,512,3)))
    model.add(LeakyReLU(0.2))
    model.add(Dropout(0.4))

    # Second Conv Block
    model.add(Conv2D(64, 5))
    model.add(LeakyReLU(0.2))
    model.add(Dropout(0.4))

    # Third Conv Block
    model.add(Conv2D(128, 5))
    model.add(LeakyReLU(0.2))
    model.add(Dropout(0.4))

    # Fourth Conv Block
    model.add(Conv2D(256, 5))
    model.add(LeakyReLU(0.2))
    model.add(Dropout(0.4))

    # Flatten then pass to dense layer
    model.add(Flatten())
    model.add(Dropout(0.4))
    model.add(Dense(1, activation='sigmoid'))

    return model 

g_opt = Adam(learning_rate=0.0001) 
d_opt = Adam(learning_rate=0.00001) 
g_loss = BinaryCrossentropy()
d_loss = BinaryCrossentropy()

class GAN(Model): 
    def __init__(self, generator, discriminator, *args, **kwargs):
        # Pass through args and kwargs to base class 
        super().__init__(*args, **kwargs)
        
        # Create attributes for gen and disc
        self.generator = generator 
        self.discriminator = discriminator 
    
    def compile(self, g_opt, d_opt, g_loss, d_loss, *args, **kwargs): 
        # Compile with base class
        super().compile(*args, **kwargs)
        
        # Create attributes for losses and optimizers
        self.g_opt = g_opt
        self.d_opt = d_opt
        self.g_loss = g_loss
        self.d_loss = d_loss 

    def train_step(self, batch):
        real_images, _ = batch
        batch_size = tf.shape(real_images)[0]
        random_latent_vectors = tf.random.normal((batch_size,512))
        fake_images = self.generator(random_latent_vectors, training=False)

        # Resize fake images to match the size of real images
        fake_images_resized = tf.image.resize(fake_images, (512, 512))
        #fake_images_resized = tf.image.grayscale_to_rgb(fake_images_resized)

        # Concatenate resized fake images with real images
        all_images = tf.concat([real_images, fake_images_resized], axis=0)

        # Labels for real and fake images
        labels_real = tf.ones((batch_size, 1))
        labels_fake = tf.zeros((batch_size, 1))
        all_labels = tf.concat([labels_real, labels_fake], axis=0)

        # Add noise to labels for discriminator
        noise_real = 0.15 * tf.random.uniform(tf.shape(labels_real))
        noise_fake = -0.15 * tf.random.uniform(tf.shape(labels_fake))
        all_labels += tf.concat([noise_real, noise_fake], axis=0)

        # Train the discriminator
        with tf.GradientTape() as d_tape:
            predictions = self.discriminator(all_images, training=True)
            d_loss = self.d_loss(all_labels, predictions)

        # Apply backpropagation for discriminator
        d_grad = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
        self.d_opt.apply_gradients(zip(d_grad, self.discriminator.trainable_variables))

        # Train the generator
        with tf.GradientTape() as g_tape:
            generated_images = self.generator(random_latent_vectors, training=True)
            predictions_fake = self.discriminator(generated_images, training=False)
            g_loss = self.g_loss(tf.ones_like(predictions_fake), predictions_fake)

        # Apply backpropagation for generator
        g_grad = g_tape.gradient(g_loss, self.generator.trainable_variables)
        self.g_opt.apply_gradients(zip(g_grad, self.generator.trainable_variables))

        return {"d_loss": d_loss, "g_loss": g_loss}

# Create instance of subclassed model
generator = build_generator()
discriminator = build_discriminator()
gan = GAN(generator, discriminator)
# Compile the model
gan.compile(g_opt, d_opt, g_loss, d_loss)

class ModelMonitor(Callback):
    def __init__(self, num_img=3, latent_dim=512):
        self.num_img = num_img
        self.latent_dim = latent_dim

    def on_epoch_end(self, epoch, logs=None):
        random_latent_vectors = tf.random.uniform((self.num_img, self.latent_dim,1))
        generated_images = self.model.generator(random_latent_vectors)
        generated_images *= 255
        generated_images.numpy()
        for i in range(self.num_img):
            img = array_to_img(generated_images[i])
            img.save(os.path.join('images', f'generated_img_{epoch}_{i}.png'))

hist = gan.fit(train_data, epochs=5, callbacks=[ModelMonitor()])

核心问题与修复方案 #

1. 修正类结构缩进错误 #

原代码中GAN类的self.generator属性、compile和train_step方法缩进错误，导致这些成员不属于类，训练时无法正确绑定模型结构，引发计算图节点缺失问题。修正后确保所有类内代码缩进正确（如上述代码所示）。

2. 移除不必要的图像resize操作 #

生成器经过5次UpSampling2D后，输出尺寸为16*2^5=512，与判别器输入(512,512,3)完全匹配，无需在训练时动态resize。移除fake_images_resized = tf.image.resize(fake_images, (512, 512))，直接使用fake_images参与后续计算：

# 替换原代码中的对应部分
all_images = tf.concat([real_images, fake_images], axis=0)

3. 修复回调中的 latent vector 形状 #

ModelMonitor中生成随机向量时多了一个维度，导致生成器输入不匹配。将：

random_latent_vectors = tf.random.uniform((self.num_img, self.latent_dim,1))

改为：

random_latent_vectors = tf.random.normal((self.num_img, self.latent_dim))

保持与训练时的向量形状和分布一致。

4. 临时绕过Grappler优化（可选） #

如果上述修复后仍触发错误，可关闭TensorFlow的重映射优化，在训练前添加：

tf.config.optimizer.set_experimental_options({"remapping": False})

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