6. 模块，层和模型

为了使用TensorFlow进行机器学习，你需要学会如何定义，保存和重载一个模型。
一个模型抽象上说有以下内容：

• 用于在张量上进行计算的函数
• 在训练过程中被不断更新的变量
  本章内容中，你会开始了解到Keras的底层，并且明白TensorFlow的模型是如何定义的，即TensorFlow如何组织变量和模型，使得模型可以被保存和重载。

模型和层在TensorFlow中的定义

大多数的模型都是有层叠加而成的。层是一种包含有可训练变量及运算的可重用的数据结构。TensorFlow中的大多数高级实现（如Keras中的层和模型）都是基于tf.Module类。
模块和层都是深度学习中的术语，用于描述有内部状态及定义这些内部状态上的操作的对象。你可以自由设置变量是否可以训练，这样方便调优。
使用tf.Module作为父类，子类将自动组织所有的tf.Variable和tf.Module对象。这使得你可以使用tf.Module中的方法来保存和重载变量。

class SimpleModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.a_variable = tf.Variable(5.0, name="train_me")
        self.non_trainable_variable = tf.Variable(5.0, trainable=False, name="do_not_train_me")

    def __call__(self, x):
        return self.a_variable * x + self.non_trainable_variable


simple_module = SimpleModule(name="simple")

simple_module(tf.constant(5.0))

print("trainable variables:", simple_module.trainable_variables)
print("all variables:", simple_module.variables)

结果是：

trainable variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>,)
all variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>, <tf.Variable 'do_not_train_me:0' shape=() dtype=float32, numpy=5.0>)

下面是一个双层神经网络的代码：

class Dense(tf.Module):
    def __init__(self, in_features, out_features, name=None):
        super().__init__(name=name)
        self.w = tf.Variable(
            tf.random.normal([in_features, out_features], name="w")
        )
        self.b = tf.Variable(
            tf.zeros([out_features], name="b")
        )

    def __call__(self, x):
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


class SequentialModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.dense_1 = Dense(in_features=3, out_features=3)
        self.dense_2 = Dense(in_features=3, out_features=2)

    def __call__(self, x):
        x = self.dense_1(x)
        return self.dense_2(x)


my_model = SequentialModule(name="the_model")

print("Model results:", my_model(tf.constant([[2., 2., 2.], [3., 3., 3.]])))

tf.Module的实现类会自动递归地收集邹游的tf.Variable和tf.Module对象。

延后创建变量

上述代码中，在模型初始化的时候便指定了输入和输出的维度，即W和b是已知维度的变量。这在大多数情况下造成了模型的局限性，很多情况下，在模型创建之前并不知道具体的维度。若能在变量第一次输入时推断出具体的输入维度，你就不需要人为指定输入维度。这样的代码更加具有灵活性。
修改之后的代码如下：

class FlexibleDenseModule(tf.Module):
    def __init__(self, out_features, name=None):
        super().__init__(name=name)
        self.is_built = False
        self.out_features = out_features

    def __call__(self, x):
        if not self.is_built:
            self.w = tf.Variable(tf.random.normal([x.shape[-1], self.out_features]), name="w")
            self.b = tf.Variable(tf.zeros([self.out_features]), name="b")
            self.is_built = True
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


class FlexibleSequentialModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.dense_1 = FlexibleDenseModule(out_features=3)
        self.dense_2 = FlexibleDenseModule(out_features=2)

    def __call__(self, x):
        x = self.dense_1(x)
        return self.dense_2(x)

变量保存

你可以将tf.Module保存为检查点（checkpoint）或存储模型（savedModel）。检查点只保存模型及子模型的变量。检查点保存时会生成两类文件：数据文件和元数据索引文件。索引文件内保存了检查点的编号，并记录了什么数据被保存了起来。数据文件记录了变量值和查找路径。你可以通过查看checkpoint的内容来确认所有变量已经保存成功。

chkp_path = "my_checkpoint"
checkpoint = tf.train.Checkpoint(model=my_model)
checkpoint.write(chkp_path)

tf.train.list_variables(chkp_path)

模型保存

TensorFlow可以在没有Python源码的情况下运行模型，这使得你可以直接从TensorFlowHub上下载已经训练好的模型来使用。TensorFlow需要在没有源码的情况下知道Python代码中描述的计算流程，我们可以使用图来达到这个目标。图记录了所有构成目标函数的计算过程。关于使用tf.function将python函数转换为图，不再赘叙。

Keras的模型和层

你可以基于tf.Module来创建属于你的高级API，而Keras也是这么做的。

Keras的层

tf.keras.lays.Layer是所有keras的层的基类，而这个基类继承自tf.Module类。
若你需要将基于tf.Module的类转换为基于keras层的类，则你只需要将父类更换一下并将call更换为call即可（keras类的call方法有自己的用处）。

class MyDense(tf.keras.layers.Layer):
    def __init__(self, in_features, out_features, **kwargs):
        super().__init__(**kwargs)

        self.w = tf.Variable(
            tf.random.normal([in_features, out_features]), name="w"
        )
        self.b = tf.Variable(
            tf.zeros([out_features]), name="b"
        )

    def call(self, x):
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


simple_layer = MyDense(name="simple", in_features=3, out_features=3)

print(simple_layer([[2., 2., 2.]]))

结果为：

tf.Tensor([[0.01792851 2.781715   0.        ]], shape=(1, 3), dtype=float32)

build阶段

前面提到过，若是在确定了输入维度之后再创建内部变量的话，是十分灵活的。
keras的层新增了一个新的生命周期（称为build）的步骤，让你更加灵活的定义和使用层。build只会被调用一次，在这个阶段确定输入的维度，多用于创建内部变量。将上面的MyDense增加build阶段，则代码如下：

class FlexibleDense(tf.keras.layers.Layer):
    def __init__(self, out_features, **kwargs):
        super().__init__(**kwargs)
        self.out_features = out_features

    def build(self, input_shape):
        self.w = tf.Variable(tf.random.normal([input_shape[-1], self.out_features]), name="w")
        self.b = tf.Variable(tf.zeros([self.out_features]), name="b")

    def call(self, x):
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


flexible_dense = FlexibleDense(out_features=3)

print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]])))

注意，由于build只会调用一次，因此若是输入与build阶段的维度不同时，会报错。

Keras模型

Keras的tf.keras.Model类提供了模型的全部特性 ，它继承自tf.keras.layers.Layer,因此keras模型可以像keras的层一样被重用和保存。 除此之外，keras的模型还提供了额外的功能，用于训练，求解，保存和重载，甚至于提供了分布式训练功能。

class FlexibleDense(tf.keras.layers.Layer):
    def __init__(self, out_features, **kwargs):
        super().__init__(**kwargs)
        self.out_features = out_features

    def build(self, input_shape):
        self.w = tf.Variable(tf.random.normal([input_shape[-1], self.out_features]), name="w")
        self.b = tf.Variable(tf.zeros([self.out_features]), name="b")

    def call(self, x):
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


class MySequentialModel(tf.keras.Model):
    def __init__(self, name=None, **kwargs):
        super().__init__(**kwargs)

        self.dense_1 = FlexibleDense(out_features=3)
        self.dense_2 = FlexibleDense(out_features=3)

    def call(self, x):
        x = self.dense_1(x)
        return self.dense_2(x)


my_sequentail_model = MySequentialModel(name="the_model")

print("Model resutls:", my_sequentail_model(tf.constant([[2., 2., 2.]])))


print(my_sequentail_model.variables)

运行结果为：

Model resutls: tf.Tensor([[0. 0. 0.]], shape=(1, 3), dtype=float32)
[<tf.Variable 'my_sequential_model/flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
array([[-1.0733198 , -2.5860493 ,  0.42328298],
       [-2.0001495 ,  1.5054438 , -0.9208656 ],
       [ 1.6782132 ,  0.72947365, -0.08435281]], dtype=float32)>, <tf.Variable 'my_sequential_model/flexible_dense/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>, <tf.Variable 'my_sequential_model/flexible_dense_1/w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 0.37761664, -1.0342877 , -0.8181074 ],
       [ 0.6091555 ,  0.97727245,  0.11385015],
       [ 1.2820089 , -0.39806262, -0.28293946]], dtype=float32)>, <tf.Variable 'my_sequential_model/flexible_dense_1/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>]