TensorFlow 自定义神经网络层实现指南

码头散步

发布时间：2026-10-13| 6人浏览过

原创

本文详细介绍如何在 TensorFlow 中自定义神经网络层。通过继承 tf.keras.Layer 类实现 __init__、build 和 call 方法，可灵活创建特定业务逻辑的层。同时，利用 tf.keras.Model 组合现有层构建复杂网络结构，如 ResNet 残差块，帮助开发者掌握模型架构设计的核心技巧。

具体实现如下所示：


import tensorflow as tf
tfe = tf.contrib.eager

tf.enable_eager_execution()

TensorFlow 图层抽象与操作

在开发机器学习模型时，通常需要在高于单个操作或变量操作的抽象层级上进行编码。

1. 关于图层的一些有用操作

许多机器学习模型可以表示为相对简单图层的组合与堆叠。TensorFlow 提供了一系列常用图层，并支持用户从头开始或作为组合来创建应用程序特定的图层。TensorFlow 在 tf.keras 包中包含了完整的 Keras API，而 Keras 图层在构建自定义模型时非常有用。


#在tf.keras.layers包中，图层是对象。要构造一个图层，只需构造一个对象。大多数层将输出维度/通道的数量作为第一个参数。
layer=tf.keras.layers.Dense(100)
#输入维度的数量通常是不必要的，因为它可以在第一次使用图层时推断出来，但如果您想手动指定它，则可以提供它，这在某些复杂模型中很有用。
layer=tf.keras.layers.Dense(10,input_shape=(None,5))
#调用层
layer(tf.zeros([10,5]))
 

#图层有许多有用的方法。例如，您可以通过调用layer.variables来检查图层中的所有变量。在这种情况下，完全连接的层将具有权重和偏差的变量。
variable=layer.variables
# variable[0]
layer.kernel.numpy()
layer.bias

自定义图层的实现方法

2. 自定义图层

实现自定义图层的最佳方式是扩展 tf.keras.Layer 类，并实现以下方法：

__init__：在此处执行所有与输入无关的初始化操作。

build 方法：在已知输入张量形状时进行剩余初始化。

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_variable("kernel", 
                  shape=[input_shape[-1].value, 
                      self.num_outputs])
  
 def call(self, input):
  return tf.matmul(input, self.kernel)
 
layer = MyDenseLayer(10)
print(layer(tf.zeros([10, 5])))
print(layer.variables)

基于 Model 的网络结构搭建

3. 搭建网络结构

机器学习模型中许多有趣的图层是通过组合现有层来实现的。例如，ResNet 中的每个残差块是卷积、批量标准化等的组合。

创建包含其他图层的类似图层时，使用的主类是 tf.keras.Model。通过继承自 tf.keras.Model 来实现一个模型。


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])
print(block(tf.zeros([1, 2, 3, 3])))
print([x.name for x in block.variables])

以上关于 TensorFlow 自定义神经网络层的实例分享完毕，希望能对读者有所帮助，也欢迎大家继续支持本站。