基于tensorflow __init__、build 和call的使用小结

基于 TensorFlow __init__、build 和 call 是一种创建自定义模型的方法。__init__ 方法通常用于初始化模型的状态（例如层权重），build 方法用于创建层权重（即，输入的形状可能未知，但输入大小会在层的第一次调用中指定），call 方法定义了前向传递逻辑。本文将详细介绍这三个方法的使用。

使用 __init__ 方法

__init__方法通常用于初始化模型状态，例如层权重、模型超参数等。在自定义层中，我们可以使用tf.keras.layer.Layer的__init__方法来进行初始化。

import tensorflow as tf

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self, input_dim, output_dim):
        super(MyDenseLayer, self).__init__()
        self.kernel = self.add_variable("kernel",
                                         shape=[input_dim, output_dim])

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

在这个示例中，我们定义一个自定义层 MyDenseLayer ，它在模型的初始化过程中，创建了一个 kernel 变量（类似于层的权重）。

使用 build 方法

build方法用于创建层权重。在调用build之前，层没有权重。在调用build之后，它至少必须创建一个权重变量（即，self.add_variable()方法），并且可以根据输入的数据来动态创建其余的权重变量。

import tensorflow as tf

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

    def build(self, input_shape):
        self.kernel = self.add_variable('kernel',
                                        shape=[input_shape[-1], self.output_dim])

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

在这个示例中，我们定义了一个自定义层MyDenseLayer。在build方法内，我们创建了一个名为kernel的权重变量，其形状为[input_shape[-1], self.output_dim]。这里，input_shape[-1]是输入张量的最后一个轴的大小（通常是输入张量的特征数量）。在后续的调用过程中，我们可以使用这个权重变量来计算前向传递逻辑。

使用 call 方法

call方法定义了层的前向传递逻辑。它的主要功能是将输入张量传递给下一层或输出。

import tensorflow as tf

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

    def build(self, input_shape):
        self.kernel = self.add_variable('kernel',
                                        shape=[input_shape[-1], self.output_dim])

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

在这个示例中，我们定义了一个自定义层 MyDenseLayer，它的前向传递计算就是输入张量与kernel权重变量的矩阵相乘。

示例1：实现一个自定义的残差网络模型

import tensorflow as tf

class ResidualBlock(tf.keras.layers.Layer):
    def __init__(self, filters, stride=1, downsample=None, name=None):
        super(ResidualBlock, self).__init__(name=name)
        self.conv1 = tf.keras.layers.Conv2D(filters=filters,
                                            kernel_size=3,
                                            strides=stride,
                                            padding="same",
                                            kernel_initializer="he_normal")
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.relu = tf.keras.layers.ReLU()
        self.conv2 = tf.keras.layers.Conv2D(filters=filters,
                                            kernel_size=3,
                                            strides=1,
                                            padding="same",
                                            kernel_initializer="he_normal")
        self.bn2 = tf.keras.layers.BatchNormalization()
        self.downsample = downsample

    def call(self, inputs, training=False):
        identity = inputs
        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x, training=training)
        if self.downsample is not None:
            identity = self.downsample(identity)
        x += identity
        x = self.relu(x)
        return x

class ResNet(tf.keras.Model):
    def __init__(self, filters, block_nums, num_classes=10):
        super(ResNet, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(filters=64,
                                            kernel_size=7,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer="he_normal")
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.maxpool = tf.keras.layers.MaxPool2D(pool_size=3,
                                                 strides=2,
                                                 padding="same")

        self.layer1 = self._make_layer(filters[0], block_nums[0])
        self.layer2 = self._make_layer(filters[1], block_nums[1], stride=2)
        self.layer3 = self._make_layer(filters[2], block_nums[2], stride=2)
        self.layer4 = self._make_layer(filters[3], block_nums[3], stride=2)

        self.avgpool = tf.keras.layers.GlobalAvgPool2D()
        self.fc = tf.keras.layers.Dense(units=num_classes, activation=tf.keras.activations.softmax)

    def _make_layer(self, filters, block_nums, stride=1):
        identity_downsample = None
        if stride != 1:
            identity_downsample = tf.keras.Sequential([
                tf.keras.layers.Conv2D(filters,
                                       kernel_size=1,
                                       strides=stride),
                tf.keras.layers.BatchNormalization(),
            ])
        res_block = []
        res_block.append(ResidualBlock(filters, stride, identity_downsample))
        for _ in range(1, block_nums):
            res_block.append(ResidualBlock(filters))

        return tf.keras.Sequential(res_block)

    def call(self, inputs, training=False):
        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.maxpool(x)

        x = self.layer1(x, training=training)
        x = self.layer2(x, training=training)
        x = self.layer3(x, training=training)
        x = self.layer4(x, training=training)

        x = self.avgpool(x)
        x = self.fc(x)

        return x

这个示例展示了如何使用自定义层和keras.Model来实现ResNet。其中自定义了残差块层 ResidualBlock，和ResNet模型层 ResNet。其中的 ResidualBlock 的实现参考了Keras官方的 ResNet50 实现。

示例2：DynamicRNN实现

import tensorflow as tf

class DynamicRNN(tf.keras.layers.Layer):
    def __init__(self, units, activation=None):
        super(DynamicRNN, self).__init__()
        self.units = units
        self.activation = tf.keras.activations.get(activation)
        self.dense_layer = tf.keras.layers.Dense(units)

    def call(self, inputs):
        _, sequence_length, _ = tf.unstack(tf.shape(inputs))
        inputs = tf.transpose(inputs, [1, 0 ,2])
        inputs = tf.reshape(inputs, [-1, tf.shape(inputs)[2]])
        outputs = self.dense_layer(inputs)
        outputs = tf.reshape(outputs, [sequence_length, -1, self.units])
        outputs = tf.transpose(outputs, [1, 0, 2])
        if self.activation is not None:
            outputs = self.activation(outputs)
        return outputs

这是一个简单的实现DynamicRNN的自定义层。DynamicRNN 层接受的输入数据是一个形如 [batch_size, sequence_length, embedding_dim] 的三维张量。其中 sequence_length 是变长的。本层的实现利用了 TensorFlow 内置函数 tf.shape 和 tf.transpose、tf.reshape 等 TensorFlow 有力的工具来实现。