Layerهایی با رفتار متفاوت در Training و Inference

گاهی Layer باید هنگام Training رفتاری متفاوت از Test داشته باشد؛ Dropout و BatchNormalization نمونه‌های استاندارد آن هستند. در این حالت آرگومان training را به call() اضافه می‌کنیم. مثال زیر در Training نویز Gaussian می‌افزاید و در Inference داده را بدون تغییر عبور می‌دهد:

class MyGaussianNoise(tf.keras.layers.Layer):
    def __init__(self, stddev, **kwargs):
        super().__init__(**kwargs)
        self.stddev = stddev

    def call(self, X, training=False):
        if training:
            noise = tf.random.normal(tf.shape(X), stddev=self.stddev)
            return X + noise
        return X

Model سفارشی با Subclassing

برای ساخت Model سفارشی از tf.keras.Model ارث‌بری می‌کنیم، Layerها را در Constructor می‌سازیم و جریان محاسبه را در call() تعریف می‌کنیم. این روش برای Architectureهایی با Loop، Branch، Skip Connection یا رفتار Dynamic مناسب است.

در مثال کتاب، ابتدا یک Layer قابل‌استفادهٔ مجدد به نام ResidualBlock ساخته می‌شود. چند Dense Layer روی ورودی اعمال می‌شوند و در پایان Output به Input اصلی اضافه می‌شود:

class ResidualBlock(tf.keras.layers.Layer):
    def __init__(self, n_layers, n_neurons, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [
            tf.keras.layers.Dense(n_neurons, activation="relu",
                                  kernel_initializer="he_normal")
            for _ in range(n_layers)]

    def call(self, inputs):
        Z = inputs
        for layer in self.hidden:
            Z = layer(Z)
        return inputs + Z

Keras به‌صورت Recursive تمام Layerهای Trackable موجود در Attributeها را شناسایی می‌کند و Variableهای آن‌ها را به Model اضافه می‌کند. سپس Model اصلی ساخته می‌شود:

class ResidualRegressor(tf.keras.Model):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden1 = tf.keras.layers.Dense(
            30, activation="relu", kernel_initializer="he_normal")
        self.block1 = ResidualBlock(2, 30)
        self.block2 = ResidualBlock(2, 30)
        self.out = tf.keras.layers.Dense(output_dim)

    def call(self, inputs):
        Z = self.hidden1(inputs)
        for _ in range(4):
            Z = self.block1(Z)
        Z = self.block2(Z)
        return self.out(Z)

Model حاصل مانند هر Model دیگر Compile، Fit، Evaluate و Predict می‌شود. اگر Serialization کامل لازم باشد، get_config() را در Classهای سفارشی پیاده‌سازی کنید؛ در غیر این صورت می‌توان فقط Weightها را Save/Load کرد.

Loss و Metric بر پایهٔ بخش‌های داخلی Model

Loss سفارشی لزوماً فقط تابع Label و Prediction نیست. گاهی می‌خواهیم Loss بر اساس Weight یا Activation داخلی Model باشد؛ برای Regularization یا Monitoring. در این حالت مقدار محاسبه‌شده را با add_loss() به Model اضافه می‌کنیم. مثال کتاب Regression DNNای با پنج Hidden Layer و یک Reconstruction Output می‌سازد تا شبکه مجبور شود اطلاعات Input را در Representationهای میانی حفظ کند.

class ReconstructingRegressor(tf.keras.Model):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [
            tf.keras.layers.Dense(30, activation="relu",
                                  kernel_initializer="he_normal")
            for _ in range(5)]
        self.out = tf.keras.layers.Dense(output_dim)
        self.reconstruction_mean = tf.keras.metrics.Mean(
            name="reconstruction_error")

    def build(self, batch_input_shape):
        n_inputs = batch_input_shape[-1]
        self.reconstruct = tf.keras.layers.Dense(n_inputs)

    def call(self, inputs, training=False):
        Z = inputs
        for layer in self.hidden:
            Z = layer(Z)
        reconstruction = self.reconstruct(Z)
        recon_loss = tf.reduce_mean(tf.square(reconstruction - inputs))
        self.add_loss(0.05 * recon_loss)
        if training:
            result = self.reconstruction_mean(recon_loss)
            self.add_metric(result)
        return self.out(Z)

ضریب ۰٫۰۵ یک Hyperparameter است تا Reconstruction Loss بر Main Loss غالب نشود. Model یا Layer می‌تواند Lossهای داخلی خود را به losses اضافه کند و Keras آن‌ها را هنگام Training جمع می‌کند.

محاسبهٔ Gradient با Automatic Differentiation

برای درک Autodiff تابع سادهٔ زیر را در نظر بگیرید:

def f(w1, w2):
    return 3 * w1 ** 2 + 2 * w1 * w2

مشتق تحلیلی نسبت به w1 برابر 6*w1 + 2*w2 و نسبت به w2 برابر 2*w1 است. در نقطهٔ (5, 3) Gradient برابر (36, 10) می‌شود. تقریب عددی با تغییر بسیار کوچک هر Parameter کار می‌کند، اما برای شبکه‌ای با هزاران یا میلیون‌ها Parameter بسیار پرهزینه است. Reverse-Mode Autodiff تمام Gradientها را با یک Forward Pass و یک Reverse Pass محاسبه می‌کند:

w1, w2 = tf.Variable(5.), tf.Variable(3.)
with tf.GradientTape() as tape:
    z = f(w1, w2)
gradients = tape.gradient(z, [w1, w2])
# [36.0, 10.0]

GradientTape Operationهای مرتبط با Variableها را ضبط می‌کند. Tape به‌طور پیش‌فرض پس از اولین فراخوانی gradient() آزاد می‌شود. اگر چند بار Gradient لازم است باید Tape را Persistent ساخت و پس از استفاده حذف کرد:

with tf.GradientTape(persistent=True) as tape:
    z = f(w1, w2)
dz_dw1 = tape.gradient(z, w1)
dz_dw2 = tape.gradient(z, w2)
del tape

به‌طور پیش‌فرض Constantها Track نمی‌شوند. برای مشتق نسبت به Tensor ثابت از tape.watch() استفاده کنید. این قابلیت برای Lossهایی که به Gradient Activation نسبت به Input وابسته‌اند مفید است.

اگر Output یک Vector باشد، gradient() Gradient مجموع مؤلفه‌ها را محاسبه می‌کند. برای Gradient جداگانهٔ هر Output نسبت به Parameterها از jacobian() استفاده می‌شود. محاسبهٔ مشتق مرتبهٔ دوم نیز ممکن است، اما در کاربردهای معمول شبکه عصبی کم‌استفاده است.

متوقف کردن Backpropagation

tf.stop_gradient() در Forward Pass مقدار Input را عبور می‌دهد، اما در Backward Pass مانند Constant عمل می‌کند:

def f(w1, w2):
    return 3 * w1 ** 2 + tf.stop_gradient(2 * w1 * w2)

در این حالت Gradient بخش دوم به w1 و w2 برنمی‌گردد.

پایداری عددی Gradientها و Custom Gradient

بعضی Functionها از نظر ریاضی درست‌اند اما در Float32 ناپایدار می‌شوند. مثلاً Gradient sqrt(x) برای x=10^-50 از محدودهٔ Float32 خارج می‌شود. راه معمول افزودن یک Epsilon کوچک است. تابع Softplus نیز اگر به شکل log(1+exp(z)) پیاده‌سازی شود برای z بزرگ Overflow می‌کند. فرم پایدارتر:

def my_softplus(z):
    return tf.math.log(1 + tf.exp(-tf.abs(z))) + tf.maximum(0., z)

گاهی خود Function پایدار است ولی Gradient تولیدشده توسط Autodiff ناپایدار است. در چنین موردی می‌توان با @tf.custom_gradient فرمول Gradient را صریحاً تعیین کرد:

@tf.custom_gradient
def my_softplus(z):
    def my_softplus_gradients(grads):
        return grads * (1 - 1 / (1 + tf.exp(z)))
    result = tf.math.log(1 + tf.exp(-tf.abs(z))) + tf.maximum(0., z)
    return result, my_softplus_gradients

تابع Gradient، Gradient رسیده از Layerهای بالاتر را می‌گیرد و طبق Chain Rule در مشتق محلی ضرب می‌کند. به این ترتیب TensorFlow اجازه می‌دهد حتی رفتار Backpropagation را برای Operation سفارشی کنترل کنیم.