Keras 中的自定义可训练层

Custom Trainable Layers in Keras

在keras中,我们可以使用Lambda层来创建自定义层,像这样:

def f(x):
    return x**2

model.add(Lambda(f))

现在我的问题是,如何让这样的自定义函数可训练?如何使此函数提高输入的幂 w,其中 w 是可训练的。像这样:

def f(x):
    return x**w

这个问题可以通过子类创建一个新层来解决,

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from tensorflow.keras.optimizers import *
import numpy as np

class ScaleLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(ScaleLayer, self).__init__()
        self.scale = tf.Variable(1., trainable=True)

    def call(self, inputs):
        return inputs ** self.scale

x = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13]).reshape(-1,1)
y = x**3.25

l = ScaleLayer()
a1 = tf.keras.layers.Input(shape=1)
a2 = l(a1)
model = tf.keras.models.Model(a1,a2)

model.compile(optimizer=Adam(learning_rate=0.01), loss='mse')
model.fit(x,y, epochs=500, verbose=0)

print(l.weights) # This prints 3.25

可以找到更多相关信息 here