如何使用 tensorflow 2 均衡学习率？

Question

我正在尝试使用 TensorFlow 版本 2 实现 StyleGAN，但我不知道如何实现均衡学习率。我尝试以这种方式缩放渐变：

def equalize_in_list(datalist):
    for i in range(len(datalist)):
        if (datalist[i] is list):
            equalize_in_list(datalist[i])
        else:
            datalist[i] = datalist[i] * np.sqrt(2)/np.prod(datalist[i].shape)
    return datalist

gen_grad = equalize_in_list(gen_grad)
disc_grad = equalize_in_list(disc_grad)

但它不能正常工作。

Answer 1

您可以只创建一个自定义图层。

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

然后像往常一样创建模型...（但您必须明确指定其参数，例如：units=x。位置参数将不起作用）

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

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

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 均衡学习率？

How may I do equalized learning rate with tensorflow 2?

python

tensorflow

generative-adversarial-network