我如何使用 tensorflow 2 进行均衡学习率? [英] How may i do equalized learning rate with tensorflow 2?

问题描述

我正在尝试使用 tensorflow 版本 2 实现 StyleGAN，但我不知道如何进行均衡学习率.我尝试用这种方式缩放渐变:

I am trying to implement StyleGAN with tensorflow version 2, and i have no ideas how to do equalized learning rate. I tried to scale gradients with this way:

但它不能正常工作.

请帮忙.

推荐答案

您可以只创建自定义图层.

You can just create a custom layer.

class DenseEQ(Dense):
    """
    Standard dense layer but includes learning rate equilization
    at runtime as per Karras et al. 2017.

    Inherits Dense layer and overides the call method.
    """
    def __init__(self, **kwargs):
        if 'kernel_initializer' in kwargs:
            raise Exception("Cannot override kernel_initializer")
        super().__init__(kernel_initializer=normal(0,1), **kwargs)

    def build(self, input_shape):
        super().build(input_shape)
        # The number of inputs
        n = np.product([int(val) for val in input_shape[1:]])
        # He initialisation constant
        self.c = np.sqrt(2/n)

    def call(self, inputs):
        output = K.dot(inputs, self.kernel*self.c) # scale kernel
        if self.use_bias:
            output = K.bias_add(output, self.bias, data_format='channels_last')
        if self.activation is not None:
            output = self.activation(output)
        return output

然后像往常一样创建模型...

And then create a model as you normally would...

model_in = Input(shape(12,))
x = DenseEq(name="whatever_1")(model_in)
x = LeakyRelu(0.2)(x)
x = DenseEq(name="whatever_2")(model_in)
model_out = LeakyRelu(0.2)(x)
model = Model(model_in, model_out)

你可以对卷积做同样的事情.

You can do the same thing for a convolution.

class Conv2DEQ(Conv2D):
    """
    Standard Conv2D layer but includes learning rate equilization
    at runtime as per Karras et al. 2017.

    Inherits Conv2D layer and overrides the call method, following
    https://github.com/keras-team/keras/blob/master/keras/layers/convolutional.py

    """
    def __init__(self, **kwargs):
        if 'kernel_initializer' in kwargs:
            raise Exception("Cannot override kernel_initializer")
        super().__init__(kernel_initializer=normal(0,1), **kwargs)

    def build(self, input_shape):
        super().build(input_shape)
        # The number of inputs
        n = np.product([int(val) for val in input_shape[1:]])
        # He initialisation constant
        self.c = np.sqrt(2/n)

    def call(self, inputs):
        if self.rank == 2:
            outputs = K.conv2d(
                inputs,
                self.kernel*self.c, # scale kernel
                strides=self.strides,
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate)

        if self.use_bias:
            outputs = K.bias_add(
                outputs,
                self.bias,
                data_format=self.data_format)

        if self.activation is not None:
            return self.activation(outputs)
        return outputs

这篇关于我如何使用 tensorflow 2 进行均衡学习率?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！