TensorFlow自定义组件开发指南分享

class NoisyLinear(tf.keras.layers.Layer):
    def __init__(self, units=32, noise_stddev=0.1):
        super().__init__()
        self.units = units
        self.noise_stddev = noise_stddev

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer="zeros",
            trainable=True
        )

    def call(self, inputs):
        noise = tf.random.normal(
            shape=tf.shape(inputs),
            stddev=self.noise_stddev
        )
        noisy_inputs = inputs + noise
        return tf.matmul(noisy_inputs, self.w) + self.b

model = tf.keras.Sequential([
    NoisyLinear(64, noise_stddev=0.2),
    tf.keras.layers.ReLU(),
    NoisyLinear(10)
])

class FocalLoss(tf.keras.losses.Loss):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def call(self, y_true, y_pred):
        ce_loss = tf.nn.sigmoid_cross_entropy_with_logits(y_true, y_pred)
        pt = tf.exp(-ce_loss)
        loss = self.alpha * tf.pow(1. - pt, self.gamma) * ce_loss
        return tf.reduce_mean(loss)

class CustomModel(tf.keras.Model):
    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred)
        
        grads = tape.gradient(loss, self.trainable_variables)
        grads, _ = tf.clip_by_global_norm(grads, 5.0)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        
        self.compiled_metrics.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

class F1Score(tf.keras.metrics.Metric):
    def __init__(self, name="f1_score"):
        super().__init__(name=name)
        self.precision = tf.keras.metrics.Precision()
        self.recall = tf.keras.metrics.Recall()

    def update_state(self, y_true, y_pred, sample_weight=None):
        self.precision.update_state(y_true, y_pred, sample_weight)
        self.recall.update_state(y_true, y_pred, sample_weight)

    def result(self):
        p = self.precision.result()
        r = self.recall.result()
        return 2 * ((p * r) / (p + r + 1e-6))

    def reset_state(self):
        self.precision.reset_state()
        self.recall.reset_state()

class L0Regularizer(tf.keras.regularizers.Regularizer):
    def __init__(self, factor=0.01):
        self.factor = factor

    def __call__(self, x):
        return self.factor * tf.reduce_sum(tf.cast(tf.not_equal(x, 0.), tf.float32))

@tf.custom_gradient
def swish(x):
    result = x * tf.nn.sigmoid(x)
    def grad(dy):
        sigmoid_x = tf.nn.sigmoid(x)
        return dy * (sigmoid_x * (1 + x * (1 - sigmoid_x)))
    return result, grad

class NoisyLinear(tf.keras.layers.Layer):
    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "noise_stddev": self.noise_stddev
        })
        return config

model = tf.keras.models.load_model(
    "model.h5",
    custom_objects={
        "NoisyLinear": NoisyLinear,
        "F1Score": F1Score
    }
)