File size: 1,932 Bytes
4483e82 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | from __future__ import annotations
import numpy as np
import tensorflow as tf
from .numpy_runtime import DEFAULT_CENTER, DEFAULT_SCALE, WINDOW_SIZE
class TensorFlowDriftModel(tf.Module):
"""TensorFlow implementation with weights shared from the PyTorch head."""
def __init__(self, linear_weight: np.ndarray, linear_bias: float) -> None:
super().__init__()
x = np.arange(WINDOW_SIZE, dtype=np.float32)
x_centered = x - x.mean()
self.x_centered = tf.constant(x_centered)
self.slope_denominator = tf.constant(np.square(x_centered).sum(), dtype=tf.float32)
self.feature_center = tf.constant(DEFAULT_CENTER)
self.feature_scale = tf.constant(DEFAULT_SCALE)
self.linear_weight = tf.constant(np.asarray(linear_weight, dtype=np.float32))
self.linear_bias = tf.constant(float(linear_bias), dtype=tf.float32)
@tf.function(
input_signature=[tf.TensorSpec([None, WINDOW_SIZE, 2], tf.float32, name="telemetry")]
)
def __call__(self, telemetry: tf.Tensor) -> dict[str, tf.Tensor]:
force = telemetry[:, :, 0]
deviation = telemetry[:, :, 1]
slope = tf.reduce_sum(force * self.x_centered, axis=1) / self.slope_denominator
shift = tf.reduce_mean(force[:, -10:], axis=1) - tf.reduce_mean(force[:, :10], axis=1)
std = tf.math.reduce_std(force, axis=1)
max_deviation = tf.reduce_max(deviation, axis=1)
last_deviation = deviation[:, -1]
force_range = tf.reduce_max(force, axis=1) - tf.reduce_min(force, axis=1)
features = tf.stack(
[slope, shift, std, max_deviation, last_deviation, force_range],
axis=1,
)
normalized = (features - self.feature_center) / self.feature_scale
logits = tf.linalg.matvec(normalized, self.linear_weight) + self.linear_bias
return {"probability": tf.math.sigmoid(logits), "features": features}
|