自定义层(06)

通过对 tf.keras.layers.Layer 进行子类化并实现以下方法来创建自定义层:

  • build:创建层的权重。使用 add_weight 方法添加权重。
  • call:定义前向传播。
  • compute_output_shape:指定在给定输入形状的情况下如何计算层的输出形状。
  • 或者,可以通过实现 get_config 方法和 from_config 类方法序列化层。

下面是一个使用核矩阵实现输入 matmul 的自定义层示例:

class MyLayer(layers.Layer):

  def __init__(self, output_dim, **kwargs):
    self.output_dim = output_dim
    super(MyLayer, self).__init__(**kwargs)

  def build(self, input_shape):
    shape = tf.TensorShape((input_shape[1], self.output_dim))
    # Create a trainable weight variable for this layer.
    self.kernel = self.add_weight(name='kernel',
                                  shape=shape,
                                  initializer='uniform',
                                  trainable=True)
    # Be sure to call this at the end
    super(MyLayer, self).build(input_shape)

  def call(self, inputs):
    return tf.matmul(inputs, self.kernel)

  def compute_output_shape(self, input_shape):
    shape = tf.TensorShape(input_shape).as_list()
    shape[-1] = self.output_dim
    return tf.TensorShape(shape)

  def get_config(self):
    base_config = super(MyLayer, self).get_config()
    base_config['output_dim'] = self.output_dim
    return base_config

  @classmethod
  def from_config(cls, config):
    return cls(**config)

使用自定义层创建模型:

model = tf.keras.Sequential([
    MyLayer(10),
    layers.Activation('softmax')])

# The compile step specifies the training configuration
model.compile(optimizer=tf.train.RMSPropOptimizer(0.001),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# Trains for 5 epochs.
model.fit(data, labels, batch_size=32, epochs=5)

Epoch 1/5
1000/1000 [==============================] - 0s 170us/step - loss: 11.4872 - acc: 0.0990
Epoch 2/5
1000/1000 [==============================] - 0s 52us/step - loss: 11.4817 - acc: 0.0910
Epoch 3/5
1000/1000 [==============================] - 0s 52us/step - loss: 11.4800 - acc: 0.0960
Epoch 4/5
1000/1000 [==============================] - 0s 57us/step - loss: 11.4778 - acc: 0.0960
Epoch 5/5
1000/1000 [==============================] - 0s 60us/step - loss: 11.4764 - acc: 0.0930

你可能感兴趣的:(自定义层(06))