TensorFlow通过子类化创建新的层和模型

通过子类化创建新的层和模型

通过本篇文章，你可以

学会通过继承tf.keras.layers.Layer类自定义层，和继承tf.keras.Model类自定神经网络模型

通过子类化创建新的层

创建步骤

# 导入相应的工具包
import tensorflow as tf
from tensorflow import keras

# 新创建的层需要继承keras.layers.Layer类
class Linear(keras.layers.Layer):
	# 在初始化函数里面定义该层所需要的权重，权重的形状、初始化、类型、是否可训练都可以定义
	def __init__(self, units=32, input_dim=32):
		super(Linear, self).__init__()  # 继承需要添加的语句super
		# 可以自定义权重的初始化
		w_init = tf.random_normal_initializer()
		# 可以自定义权重的形状，类型，是否可训练
		self.w = tf.Variable(
			initial_value = w_init(shape=(input_dim, units), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
		# 可以自定义权重的初始化
		b_init = tf.zeros_initializer()
		self.b = tf.Variable(
			initial_value = b_init(shape=( units, ), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
	# 通过call函数来定义层的运算逻辑,即前向计算逻辑
	def call(self, inputs):
		y = tf.matmul(inputs, self.w) + self.b
		return y
		
		

创建的层的权重如何提取？

linear_layer = Linear(units=8, input_dim=4)
# 获取层所有可训练的权重，结果是个列表
trainable_weights = linear.trainable_weights
# 获取层所有不可训练的权重，结果是个列表
non_trainable_weights = linear.non_trainable_weights
# 获取层所有权重，不管是否可以训练，结果是个列表
weights = linear.weights

层实例化之后再创建权重

上面的创建层的步骤需要事先知道输入的形状，其实在大多数情况中，你可能事先不知道输入的形状大小。所以我们希望可以在实例化层之后，再创建层的权重。该怎么做呢？你可以在层的build(self, inputs_shape)方法中创建层的权重，inputs_shape这个参数表示的就是输入的形状。

修改上面的例子，变成如下

# 新创建的层需要继承keras.layers.Layer类
class Linear(keras.layers.Layer):
	# 在初始化函数里面定义该层所需要的权重，权重的形状、初始化、类型、是否可训练都可以定义
	def __init__(self, units=32):
		super(Linear, self).__init__()  # 继承需要添加的语句super
		self.units = units
    def build(self, inputs_shape):
        # 可以自定义权重的初始化
		w_init = tf.random_normal_initializer()
		# 可以自定义权重的形状，类型，是否可训练
		self.w = tf.Variable(
			initial_value = w_init(shape=(inputs_shape[-1], self.units), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
		# 可以自定义权重的初始化
		b_init = tf.zeros_initializer()
		self.b = tf.Variable(
			initial_value = b_init(shape=( self.units, ), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
	# 通过call函数来定义层的运算逻辑,即前向计算逻辑
	def call(self, inputs):
		y = tf.matmul(inputs, self.w) + self.b
		return y

这是一种比较实际的做法，因为你只需创建一次实例化层，然后每次使用该层的时候，权重都会重新创建，所以权重都不一样。这是一种比较合理的做法

层可以递归组合，即层里面可以嵌套层

如果将层实例作为另一个层的组件，则外部层会跟踪内部层创建的权重

一般在外部层的初始化方法中创建子层，并在call方法中进行逻辑计算

下面创建一个外部层，以MLP模块进行示例

class MLPBlock(keras.layers.Layer):
	# 在初始化函数里面实例化子层
	def __init__(self):
        super(MLPBlock, self).__init__() 
        # 下面每个层的权重都不一样
        self.linear1 = Linear(32)  
        self.linear2 = Linear(32)
        self.linear3 = Linear(1)
    # 在call方法中进行前向逻辑计算
    def call(self, inputs):
        x = self.linear_1(inputs)
        x = tf.nn.relu(x)
        x = self