Training SimGAN, but always got mode collapse

Question

I am using SimGAN (Simulated+Unsupervised GAN) to do my task. I want to turn my synthetic data to the real image domain. But it seems not as ideal as I thought.

First I pretrained my refiner with perceptual loss to make my refiner generate my synthetic data, and it work well. here's the result and my refiner model. result

def refiner_model(width, height, channels = channels):   
    def resnet_block(input_features, nb_features=64, kernel_size=4):
        y = layers.Conv2D(nb_features, kernel_size=kernel_size, padding='same')(input_features)
        y = layers.BatchNormalization()(y)  
        y = layers.LeakyReLU(alpha=0.2)(y)  
        y = layers.Conv2D(nb_features, kernel_size=kernel_size, padding='same')(y)
        y = layers.BatchNormalization()(y)  
        y = layers.LeakyReLU(alpha=0.2)(y)  
        
        y = layers.Add()([y, input_features])
        # TODO
        y = layers.LeakyReLU(alpha=0.2)(y)
        
        
        return y

    input_layer = layers.Input(shape=(height, width, channels))
    # an input image of size w × h is convolved with 3 × 3 filters that output 64 feature maps
    x = layers.Conv2D(64, kernel_size=4, padding='same')(input_layer)
    x = layers.LeakyReLU(alpha=0.2)(x)  # 使用LeakyReLU激活函数
    for _ in range(5):
        x = resnet_block(x)
    # this will set the 
    output_layer = layers.Conv2D(channels, kernel_size=1, padding='same', activation='tanh')(x)

    return models.Model(input_layer, output_layer, name='refiner')

Then I pretrained my discriminator with my loss. I follow the paper not to use a flatten layer at the last of the discriminator, to make a local adversarial loss with BCE. Here is my discriminator and local loss.

def discriminator_model(width, height, channels):
    input_layer = layers.Input(shape=(height, width, channels))
    x = layers.Conv2D(96, (3, 3), strides=(2, 2), padding='same', kernel_initializer=initializers.HeNormal())(input_layer)
    x = layers.LeakyReLU(alpha=0.1)(x)
    x = layers.Conv2D(64, (3, 3), strides=(2, 2), padding='same')(x)
    x = layers.LeakyReLU(alpha=0.1)(x)
    x = layers.MaxPooling2D(pool_size=(3, 3), strides=(1, 1), padding='same')(x)
    x = layers.Conv2D(32, (3, 3), strides=(2, 2), padding='same')(x)
    x = layers.LeakyReLU(alpha=0.1)(x)
    x = layers.Conv2D(16, (3, 3), strides=(2, 2), padding='same')(x)
    x = layers.LeakyReLU(alpha=0.1)(x)
    x = layers.Conv2D(16, (1, 1), strides=(1, 1), padding='valid')(x)
    x = layers.LeakyReLU(alpha=0.1)(x)
    x = layers.Conv2D(1, (1, 1), strides=(1, 1), padding='valid')(x)  # Outputs local scores
    x = layers.Reshape((-1, 1))(x)
    x = layers.Activation('sigmoid')(x)  # Convert to probability

    model = models.Model(inputs=input_layer, outputs=x, name='discriminator')
    return model

def local_adversarial_loss(y_true, y_pred):
    """
    y_true shape : (batch_size, disc.output.shape[0], 1)
    y_pred shape : (batch_size, disc.output.shape[0], 1)
    Count the max loss for each image in the batch
    """
    truth = tf.reshape(y_true, (-1, 1))
    predicted = tf.reshape(y_pred, (-1, 1))
    computed_loss = tf.keras.losses.binary_crossentropy(truth , predicted , from_logits=False)
    output = tf.reduce_mean(computed_loss)
    return output

And finally I started training my SimGAN. I first trained Refiner with a combined model with untrainable discriminator. with two loss : the same local_adversarial_loss and self_regularisation_loss below.

synthetic_img = layers.Input(shape=(config.img_height, config.img_width, config.channels))
refined_output = self.refiner(synthetic_img)
discriminator_output = self.discriminator(refined_output)
self.combined_model = Model(inputs=synthetic_img, outputs=[refined_output, discriminator_output], name='combined')
self.combined_model.summary()
self.discriminator.trainable = False
self.combined_model.compile(loss={ 
                                  'refiner': my_losses.self_regularisation_loss,
                                  'discriminator': my_losses.max_local_adversarial_loss
                                  },
                            loss_weights={
                                     'refiner': 1,
                                      'discriminator': 30
                                    }, 
                             optimizer=Adam(learning_rate=config.lr_schedule_for_combined, clipvalue=0.01))
self.discriminator.trainable = True

def self_regularisation_loss(y_true, y_pred):
    # original 0.000004
    return tf.multiply(config.self_regularisation_loss, tf.reduce_sum(tf.abs(y_pred - y_true)))

Then I trained my discriminator with local_adversarial_loss too.

        for i in range(self.nb_steps):
                # train the refiner
                for _ in range(self.k_g):
                    # sample a mini-batch of synthetic images
                    syn_img_batch = utils.get_image_batch(self.syn_gen)
                    real_img_batch = utils.get_image_batch(self.real_gen)
                    # update θ by taking an SGD step on mini-batch loss LR(θ)
                    # output refine output and discriminator output
                    loss = self.combined_model.train_on_batch(syn_img_batch, [syn_img_batch, self.y_real])
                
                    self.gan_loss = np.add(self.gan_loss, loss)
                for _ in range(self.k_d):
                    # sample a mini-batch of synthetic and real images
                    syn_img_batch = utils.get_image_batch(self.syn_gen)
                    real_img_batch = utils.get_image_batch(self.real_gen)
                    
                    # refine the synthetic images w/ the current refiner
                    refined_img_batch = self.refiner.predict_on_batch(syn_img_batch)
                    
                    # use a history of refined images
                    history_img_half_batch = self.ihb.get_from_image_history_buffer()
                    self.ihb.add_to_history_img_buffer(refined_img_batch)
                    
                    if len(history_img_half_batch):
                        refined_img_batch[:self.batch_size//2] = history_img_half_batch
                    # TODO ERROR
                    # update φ by taking an SGD step on mini-batch loss LD(φ)

                    aw_method = aw_loss.AWMethod()
                    loss = aw_method.compute_aw_loss(self.discriminator, real_img_batch, syn_img_batch, my_losses.max_local_adversarial_loss)
                    # pass optimizer to the function
                    self.disc_opt.apply_gradients(loss)

I have use aw_method(a customize way to give D_real and D_fake weight when updating discriminator) and label smoothing to avoid mode collapse, but it seem still unstable and still collapse easily. Here is my real data. real data And the result is show below.

The refiner losses, the first is total loss after weighting and the second is self loss, the third is adversarial loss.

step: 2100/10000 |  loss: [76.22956827  3.04130666  2.0910932 ]
step: 2200/10000 |  loss: [72.91291349  2.93849123  1.99926921]
step: 2300/10000 |  loss: [72.77705942  2.92462765  1.99578376]
step: 2400/10000 |  loss: [59.11938229  2.69308183  1.61218002]
step: 2500/10000 |  loss: [68.23417953  3.08283158  1.86146709]
step: 2600/10000 |  loss: [83.80003647  3.51962793  2.29372596]
step: 2700/10000 |  loss: [74.07024246  3.41312608  2.01877476]
step: 2800/10000 |  loss: [82.93771881  3.79858503  2.26111811]
step: 2900/10000 |  loss: [109.05483198   3.57535326   3.01369941]
step: 3000/10000 |  loss: [80.43529324  3.37350667  2.20176533]

And the output looks not that good. output

I am wondering if I design a bad model? bade loss function even a bad training steps?