自定义层

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

我们建议使用 tf.keras 作为构建神经网络的高级 API。也就是说,大多数 TensorFlow API 都可以在 Eager Execution 模式下使用。

import tensorflow as tf
print(tf.config.list_physical_devices('GPU'))

层:常用的一组有用操作

在编写机器学习模型代码时,大多数情况下,您希望在比单个操作和单个变量操作更高的抽象级别上进行操作。

许多机器学习模型可以表示为相对简单的层的组合和堆叠,TensorFlow 提供了一组许多常用层,以及方便您从头开始或通过组合现有层来编写自己的特定于应用程序的层的方法。

TensorFlow 在 tf.keras 包中包含了完整的 Keras API,Keras 层在构建您自己的模型时非常有用。

# In the tf.keras.layers package, layers are objects. To construct a layer,
# simply construct the object. Most layers take as a first argument the number
# of output dimensions / channels.
layer = tf.keras.layers.Dense(100)
# The number of input dimensions is often unnecessary, as it can be inferred
# the first time the layer is used, but it can be provided if you want to
# specify it manually, which is useful in some complex models.
layer = tf.keras.layers.Dense(10, input_shape=(None, 5))

可以在 文档 中看到所有预先存在的层的完整列表。它包括 Dense(全连接层)、Conv2D、LSTM、BatchNormalization、Dropout 等等。

# To use a layer, simply call it.
layer(tf.zeros([10, 5]))
# Layers have many useful methods. For example, you can inspect all variables
# in a layer using `layer.variables` and trainable variables using
# `layer.trainable_variables`. In this case a fully-connected layer
# will have variables for weights and biases.
layer.variables
# The variables are also accessible through nice accessors
layer.kernel, layer.bias

实现自定义层

实现您自己的层的最佳方法是扩展 tf.keras.Layer 类并实现

  1. __init__,您可以在其中执行所有与输入无关的初始化
  2. build,您可以在其中了解输入张量的形状并执行其余的初始化
  3. call,您可以在其中执行前向计算

请注意,您不必等到调用 build 才能创建变量,您也可以在 __init__ 中创建它们。但是,在 build 中创建它们的优势在于,它允许根据层将要操作的输入的形状进行延迟变量创建。另一方面,在 __init__ 中创建变量意味着需要明确指定创建变量所需的形状。

class MyDenseLayer(tf.keras.layers.Layer):
  def __init__(self, num_outputs):
    super(MyDenseLayer, self).__init__()
    self.num_outputs = num_outputs

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])

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

layer = MyDenseLayer(10)
_ = layer(tf.zeros([10, 5])) # Calling the layer `.builds` it.
print([var.name for var in layer.trainable_variables])

如果尽可能使用标准层,代码整体上更容易阅读和维护,因为其他读者会熟悉标准层的行为。如果您想使用 tf.keras.layers 中不存在的层,请考虑提交 github 问题,或者更好的是,向我们发送一个拉取请求!

模型:组合层

机器学习模型中许多有趣的类似层的结构是通过组合现有层来实现的。例如,resnet 中的每个残差块都是卷积、批次归一化和快捷方式的组合。层可以嵌套在其他层中。

通常,当您需要模型方法(如:Model.fitModel.evaluateModel.save)时,您会从 keras.Model 继承(有关详细信息,请参阅 自定义 Keras 层和模型)。

除了跟踪变量之外,keras.Model(而不是 keras.layers.Layer)提供的另一个功能是,它还跟踪其内部层,使它们更容易检查。

例如,这是一个 ResNet 块

class ResnetIdentityBlock(tf.keras.Model):
  def __init__(self, kernel_size, filters):
    super(ResnetIdentityBlock, self).__init__(name='')
    filters1, filters2, filters3 = filters

    self.conv2a = tf.keras.layers.Conv2D(filters1, (1, 1))
    self.bn2a = tf.keras.layers.BatchNormalization()

    self.conv2b = tf.keras.layers.Conv2D(filters2, kernel_size, padding='same')
    self.bn2b = tf.keras.layers.BatchNormalization()

    self.conv2c = tf.keras.layers.Conv2D(filters3, (1, 1))
    self.bn2c = tf.keras.layers.BatchNormalization()

  def call(self, input_tensor, training=False):
    x = self.conv2a(input_tensor)
    x = self.bn2a(x, training=training)
    x = tf.nn.relu(x)

    x = self.conv2b(x)
    x = self.bn2b(x, training=training)
    x = tf.nn.relu(x)

    x = self.conv2c(x)
    x = self.bn2c(x, training=training)

    x += input_tensor
    return tf.nn.relu(x)


block = ResnetIdentityBlock(1, [1, 2, 3])
_ = block(tf.zeros([1, 2, 3, 3]))
block.layers
len(block.variables)
block.summary()

但是,大多数情况下,组合了许多层的模型只是依次调用一个层。这可以使用很少的代码使用 tf.keras.Sequential 来完成

my_seq = tf.keras.Sequential([tf.keras.layers.Conv2D(1, (1, 1),
                                                    input_shape=(
                                                        None, None, 3)),
                             tf.keras.layers.BatchNormalization(),
                             tf.keras.layers.Conv2D(2, 1,
                                                    padding='same'),
                             tf.keras.layers.BatchNormalization(),
                             tf.keras.layers.Conv2D(3, (1, 1)),
                             tf.keras.layers.BatchNormalization()])
my_seq(tf.zeros([1, 2, 3, 3]))
my_seq.summary()

下一步

现在,您可以返回到前面的笔记本,并调整线性回归示例以使用层和模型来更好地构建。