[FIXED] Correct way to apply Minibatch Standard Deviation to Keras GAN layer

Issue

I’m trying to improve the stability of my GAN model by adding a standard deviation variable to my layer’s feature map. I’m following the example set in the GANs-in-Action git. The math itself makes sense to me. The mechanics of my model and the reasons why this addresses mode collapse makes sense to me. However, a shortcoming from the example is that they never actually show how this code is executed.

def minibatch_std_layer(layer, group_size=4):
    group_size = keras.backend.minimum(group_size, tf.shape(layer)[0])

    shape = list(keras.backend.int_shape(input))
    shape[0] = tf.shape(input)[0]

    minibatch = keras.backend.reshape(layer,(group_size, -1, shape[1], shape[2], shape[3]))
    minibatch -= tf.reduce_mean(minibatch, axis=0, keepdims=True)
    minibatch = tf.reduce_mean(keras.backend.square(minibatch), axis = 0)
    minibatch = keras.backend.square(minibatch + 1e8)
    minibatch = tf.reduce_mean(minibatch, axis=[1,2,4], keepdims=True)
    minibatch = keras.backend.tile(minibatch,[group_size, 1, shape[2], shape[3]])
    return keras.backend.concatenate([layer, minibatch], axis=1)

def build_discriminator():

    const = ClipConstraint(0.01)

    discriminator_input = Input(shape=(4000,3), batch_size=BATCH_SIZE, name='discriminator_input')
    
    x = discriminator_input

    x = Conv1D(64, 3, strides=1, padding="same", kernel_constraint=const)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    x = Conv1D(128, 3, strides=2, padding="same", kernel_constraint=const)(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    x = Conv1D(256, 3, strides=3, padding="same", kernel_constraint=const)(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    # Trying to add it to the feature map here 
    x = minibatch_std_layer(Conv1D(256, 3, strides=3, padding="same", kernel_constraint=const)(x))

    x = Flatten()(x)

    x = Dense(1000)(x)

    discriminator_output = Dense(1, activation='sigmoid')(x)

    return Model(discriminator_input, discriminator_output, name='discriminator_model')

d = build_discriminator()

No matter how I structure it, I can’t get the discriminator to build. It continues to return different types of AttributeErrors but I’ve been unable to understand what it wants. Searching the issue, there were lots of Medium posts showing a high level overview of what this does in a progressive GAN, but nothing I could find showing its application.

Does anyone have any suggestions about how the above code is added to a layer?

Solution

this is my proposal…

the problem is related to the minibatch_std_layer function. first of all your network deals with 3d data while the original minibatch_std_layer deals with 4d data so you need to adapt it. secondly, the input variable defined in this function is unknown (also in the source code you cited) so I think the most obvious and logical solution is to consider it as the layer variable (the input of minibatch_std_layer). with this in mind the modified minibatch_std_layer becomes:

def minibatch_std_layer(layer, group_size=4):

    group_size = K.minimum(4, layer.shape[0])
    shape = layer.shape

    minibatch = K.reshape(layer,(group_size, -1, shape[1], shape[2]))
    minibatch -= tf.reduce_mean(minibatch, axis=0, keepdims=True)
    minibatch = tf.reduce_mean(K.square(minibatch), axis = 0)
    minibatch = K.square(minibatch + 1e-8) #epsilon=1e-8
    minibatch = tf.reduce_mean(minibatch, axis=[1,2], keepdims=True)
    minibatch = K.tile(minibatch,[group_size, 1, shape[2]])
    return K.concatenate([layer, minibatch], axis=1)

that we can put inside our model in this way:

def build_discriminator():

    # const = ClipConstraint(0.01)

    discriminator_input = Input(shape=(4000,3), batch_size=32, name='discriminator_input')
    
    x = discriminator_input

    x = Conv1D(64, 3, strides=1, padding="same")(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    x = Conv1D(128, 3, strides=2, padding="same")(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    x = Conv1D(256, 3, strides=3, padding="same")(x)
    x = LeakyReLU(0.3)(x)
    x = Dropout(0.25)(x)

    # Trying to add it to the feature map here
    x = Conv1D(256, 3, strides=3, padding="same")(x)
    x = Lambda(minibatch_std_layer)(x)

    x = Flatten()(x)

    x = Dense(1000)(x)

    discriminator_output = Dense(1, activation='sigmoid')(x)

    return Model(discriminator_input, discriminator_output, name='discriminator_model')

I don’t know what it’s ClipConstraint but It doesn’t seem problematic. I ran the code with TF 2.2 but also think that it’s quite easy to make it run with TF 1 (if u are using it). here the running code: https://colab.research.google.com/drive/1A6UNYkveuHPF7r4-XAe8MuCHZJ-1vcpl?usp=sharing

Answered By – Marco Cerliani

Answer Checked By – Robin (Easybugfix Admin)

Leave a Reply

(*) Required, Your email will not be published