Прежде всего, «унификация» API (как вы это называете) под keras не мешает вам делать то, что вы делали в TensorFlow 1.x. Сеансы могут быть пропущены, но вы все равно можете определить свою модель как любую функцию Python и обучать ее с нетерпением без keras (то есть до tf.GradientTape )
Теперь, если вы хотите построить модель keras с пользовательским слоем , который выполняет пользовательскую операцию и имеет пользовательский градиент , вы должны выполнить следующее:
a) Напишите функцию, которая выполняет вашу пользовательскую операцию и определит ваш пользовательский градиент. Подробнее о том, как это сделать здесь .
@tf.custom_gradient
def custom_op(x):
result = ... # do forward computation
def custom_grad(dy):
grad = ... # compute gradient
return grad
return result, custom_grad
Обратите внимание, что в функции вы должны трактовать x и dy как Tensors, а не numpy массивы (т.е. выполнять тензорные операции)
b) Создайте пользовательский слой keras, который выполняет ваш custom_op. В этом примере я предполагаю, что ваш слой не имеет обучаемых параметров или изменяет форму ввода, но это не имеет большого значения, если он есть. Для этого вы можете обратиться к руководству, которое вы разместили, отметьте этот .
class CustomLayer(tf.keras.layers.Layer):
def __init__(self):
super(CustomLayer, self).__init__()
def call(self, x):
return custom_op(x) # you don't need to explicitly define the custom gradient
# as long as you registered it with the previous method
Теперь вы можете использовать этот слой в модели керас, и он будет работать. Например:
inp = tf.keras.layers.Input(input_shape)
conv = tf.keras.layers.Conv2D(...)(inp) # add params like the number of filters
cust = CustomLayer()(conv) # no parameters in custom layer
flat = tf.keras.layers.Flatten()(cust)
fc = tf.keras.layers.Dense(num_classes)(flat)
model = tf.keras.models.Model(inputs=[inp], outputs=[fc])
model.compile(loss=..., optimizer=...) # add loss function and optimizer
model.fit(...) # fit the model