TypeError: использование пользовательской функции активации при включенной смешанной точности?
Я пытался использовать пользовательскую активацию в обучающих конвейерах со смешанной точностью, но столкнулся со следующей ошибкой:
TypeError: Input 'y' of 'Mul' Op has type float32 that does not match type float16 of argument 'x'.
Воспроизвести
Включение смешанной точности...
import tensorflow as tf
policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)
print('Mixed precision enabled')
Пользовательская активация...
def ARelu(x, alpha=0.90, beta=2.0):
alpha = tf.clip_by_value(alpha, clip_value_min=0.01, clip_value_max=0.99)
beta = 1 + tf.math.sigmoid(beta)
return tf.nn.relu(x) * beta - tf.nn.relu(-x) * alpha
Обучение...
import tensorflow as tf
(xtrain, ytrain), (xtest, ytest) = tf.keras.datasets.mnist.load_data()
def pre_process(inputs, targets):
inputs = tf.expand_dims(inputs, -1)
targets = tf.one_hot(targets, depth=10)
return tf.divide(inputs, 255), targets
train_data = tf.data.Dataset.from_tensor_slices((xtrain, ytrain)).\
take(10_000).shuffle(10_000).batch(8).map(pre_process)
test_data = tf.data.Dataset.from_tensor_slices((xtest, ytest)).\
take(1_000).shuffle(1_000).batch(8).map(pre_process)
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), strides=(1, 1),
input_shape=(28, 28, 1), activation=ARelu),
tf.keras.layers.MaxPool2D(pool_size=(2, 2)),
tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1),
activation=ARelu),
tf.keras.layers.MaxPool2D(pool_size=(2, 2)),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation=ARelu),
tf.keras.layers.Dense(10, activation='softmax', dtype=tf.float32)])
opt = tf.keras.optimizers.Adam()
model.compile(loss='categorical_crossentropy', optimizer=opt)
history = model.fit(train_data, validation_data=test_data, epochs=10)
# ------------------
TypeError: Input 'y' of 'Mul' Op has type float32 that does not match type float16 of argument 'x'.
Однако без смешанной точности это работает. Я понимаю, что проблема просто вводит пропущенное совпадение, но где я должен это изучить?
Кроме того, пытаясь решить эту проблему, я обнаружил, что использование tf.keras.mixed_precision.LossScaleOptimizer безопасно, чтобы избежать числового недополнения. Это что-то, что мы должны использовать для обучения смешанной точности?
Ответы 2
Чтобы решить эту проблему, я должен передать ввод float32. Я не уверен, что это правильный способ решить эту ошибку.
def ARelu(x, alpha=0.90, beta=2.0):
alpha = tf.clip_by_value(alpha, clip_value_min=0.01, clip_value_max=0.99)
beta = 1 + tf.math.sigmoid(beta)
x = tf.cast(x, 'float32')
return tf.nn.relu(x) * beta - tf.nn.relu(-x) * alpha
Просто приведя тип к float32, это работает.
Подробности
Однако дело в том, что для того, чтобы использовать mixed-precision, мы должны сделать следующее:
# At the beginning ....
policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)
print('Mixed precision enabled')
и
# at the last layer, adding dtype as float 32
tf.keras.layers.Dense(num_classes, activation=..., dtype=tf.float32)])
Честно говоря, до сих пор не знаю, как работает механизм mixed-precision сзади. Сначала он устанавливает политику mixed_float16 и активирует вывод как tf.float32. Таким образом, мы не могли использовать пользовательскую функцию активации, если только не преобразовали ввод x в float32, что, как я полагаю, было float16 для смешанной точности.
Решение вышеуказанной проблемы заключается в преобразовании ваших определенных альфа- и бета-каналов в float16, а не в преобразовании входных данных вашего слоя активации в Float32.
ПОДРОБНОСТИ:
На самом деле, основная причина использования МП заключается в уменьшении объема памяти, наблюдаемого во время обучения. Метод для этого заключается в сохранении вывода слоя в FP16, поскольку в потреблении памяти преобладает хранение активаций, а не весов. Преобразовывая выходные данные слоя в FP32 в пользовательской функции активации, вы теряете эту экономию и даже требуете больше памяти для обучения модели по сравнению с использованием полной точности, потому что для вашей активации существует 2 копии.
Что касается преобразования вашего вывода в float32, это делается только для последнего слоя. В вашем случае вы переводите все выходные слои (следующий ввод слоя) в float32 с помощью x = tf.cast(x, 'float32'), что не рекомендуется. Что касается нескольких копий, поскольку tf.Gradienttape, скорее всего, используется для выполнения обратного распространения, все операции и тензоры будут сохранены, чтобы упростить вычисление ваших градиентов. Это часто приводит к сохранению обеих точности ваших тензоров до тех пор, пока обратный путь не будет завершен. Надеюсь это поможет!
Просто попробуйте привести альфа и бета к половинной точности, а не ввести вход x, и проверьте пиковую память, необходимую для одного шага обучения, используя tf.print('PEAK MEMORY in GB: ', tf.config.experimental.get_memory_info('GPU:0')['peak'] / 1000000000)
В наши дни обычной практикой является преобразование в float32 из последнего слоя. Тем не менее, не могли бы вы подробнее рассказать о ваших последних утверждениях с исходным кодом относительно существования нескольких копий? Спасибо.