在 TensorFlow 中构建自定义优化器-Python教程-PHP中文网

在 TensorFlow 中构建自定义优化器

DDD

发布： 2025-10-02 19:05:00

原创

246人浏览过

在 tensorflow 中构建自定义优化器

本文档旨在指导开发者如何在 TensorFlow 中创建自定义优化算法。我们将深入探讨如何获取每次迭代的当前点向量 x 和梯度向量 g，以及如何更新 x 并将更新后的值设置回模型。通过一个具体的示例，我们将展示如何修改梯度形状以适应自定义优化算法的需求，并提供构建和应用自定义优化器的完整流程。

自定义优化器类

在 TensorFlow 中创建自定义优化器，需要继承 tf.keras.optimizers.Optimizer 类，并重写其关键方法。以下是一个基本框架：

from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_training_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.training import optimizer
from tensorflow.python.util.tf_export import tf_export
import tensorflow as tf
import numpy as np

class CustomOptimizer(optimizer.Optimizer):
  def __init__(self, learning_rate=0.01, use_locking=False, name="CustomOptimizer"):
    super(CustomOptimizer, self).__init__(use_locking, name)
    self._learning_rate = learning_rate

  def _create_slots(self, var_list):
    # 初始化优化器所需的变量，例如动量、学习率等。
    pass

  def _prepare(self):
    # 将学习率等参数转换为 Tensor。
    self._learning_rate_tensor = ops.convert_to_tensor(self._learning_rate, name="learning_rate")

  def _apply_dense(self, grad, var):
    # 对稠密张量应用梯度更新。这是核心逻辑所在。
    return self._resource_apply_dense(grad, var)

  def _resource_apply_dense(self, grad, var):
    # 使用资源变量应用梯度更新。
    var_update = tf.compat.v1.assign_sub(var, self._learning_rate_tensor * grad)
    return tf.group(var_update)

  def _apply_sparse(self, grad, var):
    # 对稀疏张量应用梯度更新。通常需要抛出 NotImplementedError。
    raise NotImplementedError("Sparse gradient updates are not supported.")

登录后复制

关键方法详解

__init__(self, ...): 构造函数，用于初始化优化器的参数，如学习率、动量等。
_create_slots(self, var_list): 创建优化器需要的辅助变量（slots），例如动量累积量。var_list 是需要优化的变量列表。
_prepare(self): 将 Python 数值类型的超参数转换为 TensorFlow 的 Tensor 类型，方便在计算图中使用。
_apply_dense(self, grad, var): 对稠密梯度进行更新。grad 是梯度 Tensor，var 是需要更新的变量 Tensor。
_resource_apply_dense(self, grad, var): 使用资源变量进行梯度更新。这是实际执行更新操作的地方。
_apply_sparse(self, grad, var): 对稀疏梯度进行更新。如果你的优化器不支持稀疏梯度，可以抛出 NotImplementedError。

获取梯度和变量

在 _apply_dense 或 _resource_apply_dense 方法中，你可以访问到当前迭代的梯度 grad 和变量 var。然而，它们的形状可能并不是你期望的向量形式。通常，grad 的形状会与变量 var 的形状一致，例如卷积层的权重矩阵。

为了将 grad 转换为向量，可以使用 tf.reshape 函数将其扁平化：

PHPEIP

PhpEIP企业信息化平台主要解决企业各类信息的集成，能把各种应用系统(如内容管理系统，网上商城，论坛系统等)统一到企业信息化平台中，整个系统采用简单易用的模板引擎，可自定义XML标签，系统采用开放式模块开发，符合开发接口的模块可完全嵌入到平台;内容管理模块可自定义内容模型，系统自带普通文章模型和图片集模型，用户可以定义丰富的栏目构建企业门户，全站可生成静态页面，提供良好的搜索引擎优化;会员管理模

查看详情

def _apply_dense(self, grad, var):
    # 将梯度扁平化为 1D 向量
    grad_flat = tf.reshape(grad, [-1])
    # 使用扁平化的梯度进行更新
    var_update = self._resource_apply_dense(grad_flat, var)
    return tf.group(var_update)

登录后复制

更新变量

使用 tf.compat.v1.assign_sub 函数可以更新变量的值。例如，要实现简单的梯度下降，可以这样更新变量：

def _resource_apply_dense(self, grad, var):
    # 使用学习率更新变量
    var_update = tf.compat.v1.assign_sub(var, self._learning_rate_tensor * grad)
    return tf.group(var_update)

登录后复制

示例：LeNet-5 模型与自定义优化器

以下是一个完整的示例，展示如何使用自定义优化器训练 LeNet-5 模型：

import tensorflow as tf

# 自定义优化器
class CustomOptimizer(tf.keras.optimizers.Optimizer):
  def __init__(self, learning_rate=0.01, name="CustomOptimizer"):
    super().__init__(name=name)
    self.learning_rate = learning_rate

  def _create_slots(self, var_list):
      pass

  def _resource_apply_dense(self, grad, var):
      var_update = var.assign_sub(self.learning_rate * grad)
      return tf.group(var_update)

  def get_config(self):
        config = super().get_config()
        config.update({
            "learning_rate": self.learning_rate,
        })
        return config

# 构建 LeNet-5 模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(6, kernel_size=(5, 5), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Conv2D(16, kernel_size=(5, 5), activation='relu'),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(120, activation='relu'),
    tf.keras.layers.Dense(84, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 使用自定义优化器
custom_optimizer = CustomOptimizer(learning_rate=0.001)

# 编译模型
model.compile(optimizer=custom_optimizer,
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 加载 MNIST 数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

x_train = x_train[..., tf.newaxis].astype("float32")
x_test = x_test[..., tf.newaxis].astype("float32")

# 创建数据集
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(60000).batch(64)
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(64)

# 训练模型
model.fit(train_dataset, epochs=5)

# 评估模型
test_loss, test_acc = model.evaluate(test_dataset)
print(f"Test accuracy: {test_acc}")

登录后复制