Implementation of Equalized Learning Rate in Keras

Question

I'm trying to implement an equalized learning rate in my GAN. I've initialized weights with normal distribution with stddev and changed a call function of my layer to:

def call(self, inputs):
    he_constant = tf.sqrt(
        x=2. / tf.size(input=self.kernel, out_type=tf.float32)
    )
    self.kernel = self.kernel * he_constant
    return super(EqualizedConv2D, self).call(inputs)

But calculated he_constant is extremely small, e.g. 0.004 and this results in network not learning anything. I've ended by artificially increasing this value

What am I missing?

Answer 1

I guess, that I've found the solution, I'm using scaled weights, not replacing them. As a result, GAN learns nicely and loss is more reasonable. New implementation:

def call(self, inputs):
    return super(EqualizedConv2D, self).call(inputs)
    # --- disabled rest equalized learning rate for now, does not work as expected.
    outputs = backend.conv2d(
        inputs,
        self.kernel*self.scale,
        strides=self.strides,
        padding=self.padding,
        data_format=self.data_format,
        dilation_rate=self.dilation_rate)
    outputs = backend.bias_add(
        outputs,
        self.bias,
        data_format=self.data_format)
    if self.activation is not None:
        return self.activation(outputs)
    return outputs

If it's still not correct, please let me know.