Класс LossScaleOptimizer

Класс LossScaleOptimizer применяется для оборачивания обычного оптимизатора с целью динамического масштабирования значений потерь. Это нужно при обучении в смешанной точности mixed_float16, где градиенты могут становиться слишком маленькими и терять точность из-за переполнения в формате float16. Первым параметром передается исходный оптимизатор, например Adam или SGD, а вторым - объект политики смешанной точности mixed_precision.Policy. Масштабирование потерь увеличивает их перед вычислением градиентов, а затем градиенты делятся обратно на тот же коэффициент перед применением к переменным.

Синтаксис

tf.keras.mixed_precision.LossScaleOptimizer(
	inner_optimizer,
	loss_scale="dynamic"
)

Пример

Давайте создадим обертку вокруг оптимизатора Adam с динамическим масштабированием потерь:

import tensorflow as tf

inner = tf.keras.optimizers.Adam(learning_rate=0.001)
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(inner)

print(optimizer)

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

<keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer object at 0x...>

Пример

Давайте посмотрим, как получить текущее значение масштаба потерь и внутренний оптимизатор:

import tensorflow as tf

inner = tf.keras.optimizers.SGD(learning_rate=0.01)
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(inner)

print(optimizer.loss_scale)
print(optimizer.inner_optimizer)

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

tf.Tensor(32768.0, shape=(), dtype=float32)
<keras.optimizers.sgd.SGD object at 0x...>

Пример

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

import tensorflow as tf

tf.random.set_seed(0)
tf.keras.mixed_precision.set_global_policy("mixed_float16")

model = tf.keras.Sequential([
	tf.keras.layers.Dense(4, activation="relu"),
	tf.keras.layers.Dense(1)
])

inner = tf.keras.optimizers.Adam(learning_rate=0.001)
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(inner)

loss_fn = tf.keras.losses.MeanSquaredError()
x = tf.constant([[1.0], [2.0], [3.0], [4.0]])
y = tf.constant([[2.0], [4.0], [6.0], [8.0]])

with tf.GradientTape() as tape:
	predictions = model(x, training=True)
	loss = loss_fn(y, predictions)
	scaled_loss = optimizer.get_scaled_loss(loss)

scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
gradients = optimizer.get_unscaled_gradients(scaled_gradients)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

print(loss.numpy())

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

...

Пример

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

import tensorflow as tf

tf.random.set_seed(0)
tf.keras.mixed_precision.set_global_policy("mixed_float16")

model = tf.keras.Sequential([
	tf.keras.layers.Dense(4, activation="relu"),
	tf.keras.layers.Dense(1)
])

inner = tf.keras.optimizers.SGD(learning_rate=0.01)
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(inner)

loss_fn = tf.keras.losses.MeanSquaredError()
x = tf.constant([[1.0], [2.0], [3.0], [4.0]])
y = tf.constant([[2.0], [4.0], [6.0], [8.0]])

def train_step():
	with tf.GradientTape() as tape:
		predictions = model(x, training=True)
		loss = loss_fn(y, predictions)
	optimizer.minimize(loss, model.trainable_variables, tape=tape)
	return loss

for step in range(3):
	loss = train_step()
	print("Шаг", step, "потери:", loss.numpy())

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

...

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

функцию set_global_policy,
которая задает глобальную политику смешанной точности
функцию global_policy,
которая возвращает текущую глобальную политику
класс Policy,
который описывает политику вычислений и переменных
функцию custom_gradient,
которая позволяет задавать собственный градиент