Метод update_step класса Optimizer

Метод update_step класса Optimizer выполняет один шаг оптимизации: применяет градиенты к переменным модели и обновляет их значения. Метод принимает градиенты и переменные, а также может принимать дополнительные аргументы, такие как имя операции. Он используется внутри метода apply_gradients и редко вызывается напрямую, но понимание его работы помогает при создании собственных оптимизаторов.

Синтаксис

optimizer.update_step(grads, vars, [name])

Пример

Давайте создадим простой оптимизатор и переменную, вычислим градиент, а затем выполним шаг обновления вручную, чтобы увидеть, как изменяется значение переменной:

import tensorflow as tf

tf.random.set_seed(0)

optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
var = tf.Variable([1.0, 2.0, 3.0])

with tf.GradientTape() as tape:
	loss = tf.reduce_sum(var ** 2)

grads = tape.gradient(loss, var)
print("Gradients:", grads)

optimizer.update_step([grads], [var])
print("Updated variable:", var.numpy())

Результат выполнения кода:

Gradients: tf.Tensor([2. 4. 6.], shape=(3,), dtype=float32)
Updated variable: [0.8 1.6 2.4]

Пример

Давайте посмотрим, как метод update_step используется внутри метода apply_gradients на примере двух переменных:

import tensorflow as tf

tf.random.set_seed(0)

optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
var1 = tf.Variable([1.0, 2.0])
var2 = tf.Variable([3.0, 4.0])

with tf.GradientTape() as tape:
	loss = tf.reduce_sum(var1 ** 2) + tf.reduce_sum(var2 ** 2)

grads = tape.gradient(loss, [var1, var2])
print("Gradients for var1:", grads[0])
print("Gradients for var2:", grads[1])

optimizer.apply_gradients(zip(grads, [var1, var2]))
print("Updated var1:", var1.numpy())
print("Updated var2:", var2.numpy())

Результат выполнения кода:

Gradients for var1: tf.Tensor([2. 4.], shape=(2,), dtype=float32)
Gradients for var2: tf.Tensor([6. 8.], shape=(2,), dtype=float32)
Updated var1: [0.8 1.6]
Updated var2: [2.4 3.2]

Пример

Давайте создадим собственный оптимизатор, который переопределяет метод update_step для реализации простого градиентного спуска:

import tensorflow as tf

tf.random.set_seed(0)

class MySGD(tf.keras.optimizers.Optimizer):
	def __init__(self, learning_rate=0.01, name="MySGD", **kwargs):
		super().__init__(name=name, **kwargs)
		self._learning_rate = learning_rate
	
	def update_step(self, gradient, variable, learning_rate):
		variable.assign_sub(learning_rate * gradient)
	
	def get_config(self):
		config = super().get_config()
		config.update({"learning_rate": self._learning_rate})
		return config

optimizer = MySGD(learning_rate=0.1)
var = tf.Variable([1.0, 2.0, 3.0])

with tf.GradientTape() as tape:
	loss = tf.reduce_sum(var ** 2)

grads = tape.gradient(loss, var)
print("Gradients:", grads)

optimizer.apply_gradients([(grads, var)])
print("Updated variable:", var.numpy())

Результат выполнения кода:

Gradients: tf.Tensor([2. 4. 6.], shape=(3,), dtype=float32)
Updated variable: [0.8 1.6 2.4]

Смотрите также

класс Optimizer,
который является базовым классом для всех оптимизаторов
метод apply_gradients,
который применяет градиенты к переменным
метод minimize,
который минимизирует функцию потерь
атрибут learning_rate,
который задает скорость обучения