使用 TensorFlow NumPy 编写 Keras 模型

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import tensorflow.experimental.numpy as tnp
import keras
from keras import layers

(x_train, y_train), (x_test, y_test) = keras.datasets.boston_housing.load_data(
    path="boston_housing.npz", test_split=0.2, seed=113
)
input_dim = x_train.shape[1]


def evaluate_model(model: keras.Model):
    loss, percent_error = model.evaluate(x_test, y_test, verbose=0)
    print("训练前的平均绝对百分比误差: ", percent_error)
    model.fit(x_train, y_train, epochs=200, verbose=0)
    loss, percent_error = model.evaluate(x_test, y_test, verbose=0)
    print("训练后的平均绝对百分比误差:", percent_error)

class TNPForwardFeedRegressionNetwork(keras.Model):
    def __init__(self, blocks=None, **kwargs):
        super().__init__(**kwargs)
        if not isinstance(blocks, list):
            raise ValueError(f"blocks 必须是一个列表，获取到的 blocks={blocks}")
        self.blocks = blocks
        self.block_weights = None
        self.biases = None

    def build(self, input_shape):
        current_shape = input_shape[1]
        self.block_weights = []
        self.biases = []
        for i, block in enumerate(self.blocks):
            self.block_weights.append(
                self.add_weight(
                    shape=(current_shape, block),
                    trainable=True,
                    name=f"block-{i}",
                    initializer="glorot_normal",
                )
            )
            self.biases.append(
                self.add_weight(
                    shape=(block,),
                    trainable=True,
                    name=f"bias-{i}",
                    initializer="zeros",
                )
            )
            current_shape = block

        self.linear_layer = self.add_weight(
            shape=(current_shape, 1),
            name="linear_projector",
            trainable=True,
            initializer="glorot_normal",
        )

    def call(self, inputs):
        activations = inputs
        for w, b in zip(self.block_weights, self.biases):
            activations = tnp.matmul(activations, w) + b
            # ReLu 激活函数
            activations = tnp.maximum(activations, 0.0)

        return tnp.matmul(activations, self.linear_layer)

model = TNPForwardFeedRegressionNetwork(blocks=[3, 3])
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
evaluate_model(model)

def tnp_mse(y_true, y_pred):
    return tnp.mean(tnp.square(y_true - y_pred), axis=0)


keras.backend.clear_session()
model = TNPForwardFeedRegressionNetwork(blocks=[3, 3])
model.compile(
    optimizer="adam",
    loss=tnp_mse,
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
evaluate_model(model)

def tnp_relu(x):
    return tnp.maximum(x, 0)


class TNPDense(keras.layers.Layer):
    def __init__(self, units, activation=None):
        super().__init__()
        self.units = units
        self.activation = activation

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

    def call(self, inputs):
        outputs = tnp.matmul(inputs, self.w) + self.bias
        if self.activation:
            return self.activation(outputs)
        return outputs


def create_layered_tnp_model():
    return keras.Sequential(
        [
            TNPDense(3, activation=tnp_relu),
            TNPDense(3, activation=tnp_relu),
            TNPDense(1),
        ]
    )


model = create_layered_tnp_model()
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
model.build((None, input_dim))
model.summary()

evaluate_model(model)

def create_mixed_model():
    return keras.Sequential(
        [
            TNPDense(3, activation=tnp_relu),
            # 模型将可以无缝使用普通的Dense层
            layers.Dense(3, activation="relu"),
            # ... 或者切换回tnp层！
            TNPDense(1),
        ]
    )


model = create_mixed_model()
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
model.build((None, input_dim))
model.summary()

evaluate_model(model)

gpus = tf.config.list_logical_devices("GPU")
if gpus:
    strategy = tf.distribute.MirroredStrategy(gpus)
else:
    # 我们可以回退到一个无操作的 CPU 策略。
    strategy = tf.distribute.get_strategy()
print("Running with strategy:", str(strategy.__class__.__name__))

with strategy.scope():
    model = create_layered_tnp_model()
    model.compile(
        optimizer="adam",
        loss="mean_squared_error",
        metrics=[keras.metrics.MeanAbsolutePercentageError()],
    )
    model.build((None, input_dim))
    model.summary()
    evaluate_model(model)

keras.backend.clear_session()

%load_ext tensorboard

models = [
    (
        TNPForwardFeedRegressionNetwork(blocks=[3, 3]),
        "TNPForwardFeedRegressionNetwork",
    ),
    (create_layered_tnp_model(), "layered_tnp_model"),
    (create_mixed_model(), "mixed_model"),
]
for model, model_name in models:
    model.compile(
        optimizer="adam",
        loss="mean_squared_error",
        metrics=[keras.metrics.MeanAbsolutePercentageError()],
    )
    model.fit(
        x_train,
        y_train,
        epochs=200,
        verbose=0,
        callbacks=[keras.callbacks.TensorBoard(log_dir=f"logs/{model_name}")],
    )

%tensorboard --logdir logs