代码示例 / Keras 速查手册 / 使用 TensorFlow NumPy 编写 Keras 模型

使用 TensorFlow NumPy 编写 Keras 模型

作者: lukewood
创建日期 2021/08/28
上次修改 2021/08/28
描述:概述如何使用 TensorFlow NumPy API 编写 Keras 模型。

ⓘ 此示例使用 Keras 3

在 Colab 中查看 GitHub 源代码


简介

NumPy 是一个非常成功的 Python 线性代数库。

TensorFlow 最近推出了 tf_numpy,它是 NumPy API 的一个大型子集的 TensorFlow 实现。感谢 tf_numpy,您可以使用 NumPy 风格编写 Keras 层或模型!

TensorFlow NumPy API 与 TensorFlow 生态系统完全集成。自动微分、TensorBoard、Keras 模型回调、TPU 分布和模型导出等功能都得到支持。

让我们浏览一些示例。


设置

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("Mean absolute percent error before training: ", percent_error)
    model.fit(x_train, y_train, epochs=200, verbose=0)
    loss, percent_error = model.evaluate(x_test, y_test, verbose=0)
    print("Mean absolute percent error after training:", percent_error)
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/california_housing.npz
 743530/743530 ━━━━━━━━━━━━━━━━━━━━ 0s 0us/step

使用 TNP 对 keras.Model 进行子类化

利用 Keras API 的最灵活方法是对 [keras.Model](/api/models/model#model-class) 类进行子类化。对 Model 类进行子类化使您可以完全自定义训练循环中发生的事情。这使得对 Model 进行子类化成为研究人员的常用选择。

在此示例中,我们将实现一个 Model 子类,该子类使用 TNP API 对波士顿房价数据集执行回归。请注意,在将 TNP API 与 keras 结合使用时,会自动处理微分和梯度下降。

首先,让我们定义一个简单的 TNPForwardFeedRegressionNetwork 类。

class TNPForwardFeedRegressionNetwork(keras.Model):
    def __init__(self, blocks=None, **kwargs):
        super().__init__(**kwargs)
        if not isinstance(blocks, list):
            raise ValueError(f"blocks must be a list, got 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 activation function
            activations = tnp.maximum(activations, 0.0)

        return tnp.matmul(activations, self.linear_layer)

就像任何其他 Keras 模型一样,我们可以使用我们想要的任何受支持的优化器、损失、指标或回调。

让我们看看模型的表现如何!

model = TNPForwardFeedRegressionNetwork(blocks=[3, 3])
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
evaluate_model(model)
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1699909864.025985   48611 device_compiler.h:187] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.

Mean absolute percent error before training:  99.96772766113281
Mean absolute percent error after training: 40.94866180419922

太好了!我们的模型似乎正在有效地学习解决手头的问题。

我们还可以使用 TNP 编写我们自己的自定义损失函数。

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)
Mean absolute percent error before training:  99.99896240234375
Mean absolute percent error after training: 52.533199310302734

使用 TNP 实现基于 Keras 层的模型

如果需要,TNP 也可以用于面向层的 Keras 代码结构。让我们实现同一个模型,但使用分层的方法!

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)
Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ tnp_dense (TNPDense)            │ (None, 3)                 │         27 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ tnp_dense_1 (TNPDense)          │ (None, 3)                 │         12 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ tnp_dense_2 (TNPDense)          │ (None, 1)                 │          4 │
└─────────────────────────────────┴───────────────────────────┴────────────┘
 Total params: 43 (172.00 B)
 Trainable params: 43 (172.00 B)
 Non-trainable params: 0 (0.00 B)
Mean absolute percent error before training:  100.00006866455078
Mean absolute percent error after training: 43.57806396484375

您还可以无缝地在 TNP 层和原生 Keras 层之间切换!

def create_mixed_model():
    return keras.Sequential(
        [
            TNPDense(3, activation=tnp_relu),
            # The model will have no issue using a normal Dense layer
            layers.Dense(3, activation="relu"),
            # ... or switching back to tnp layers!
            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)
Model: "sequential_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ tnp_dense_3 (TNPDense)          │ (None, 3)                 │         27 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ dense (Dense)                   │ (None, 3)                 │         12 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ tnp_dense_4 (TNPDense)          │ (None, 1)                 │          4 │
└─────────────────────────────────┴───────────────────────────┴────────────┘
 Total params: 43 (172.00 B)
 Trainable params: 43 (172.00 B)
 Non-trainable params: 0 (0.00 B)
Mean absolute percent error before training:  100.0
Mean absolute percent error after training: 55.646610260009766

Keras API 提供了各种各样的层。能够将它们与 NumPy 代码一起使用可以极大地节省项目中的时间。


分布策略

TensorFlow NumPy 和 Keras 与 TensorFlow 分布式策略 集成。这使得在多个 GPU 甚至整个 TPU Pod 上执行分布式训练变得简单。

gpus = tf.config.list_logical_devices("GPU")
if gpus:
    strategy = tf.distribute.MirroredStrategy(gpus)
else:
    # We can fallback to a no-op CPU strategy.
    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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
Running with strategy: MirroredStrategy
Model: "sequential_2"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ tnp_dense_5 (TNPDense)          │ (None, 3)                 │         27 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ tnp_dense_6 (TNPDense)          │ (None, 3)                 │         12 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ tnp_dense_7 (TNPDense)          │ (None, 1)                 │          4 │
└─────────────────────────────────┴───────────────────────────┴────────────┘
 Total params: 43 (172.00 B)
 Trainable params: 43 (172.00 B)
 Non-trainable params: 0 (0.00 B)
Mean absolute percent error before training:  100.0
Mean absolute percent error after training: 44.573463439941406

TensorBoard 集成

使用 Keras API 的众多好处之一是能够通过 TensorBoard 监控训练。将 TensorFlow NumPy API 与 Keras 一起使用使您可以轻松利用 TensorBoard。

keras.backend.clear_session()

要从 Jupyter notebook 加载 TensorBoard,您可以运行以下魔法命令

%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}")],
    )
/opt/conda/envs/keras-tensorflow/lib/python3.10/site-packages/keras/src/callbacks/tensorboard.py:676: UserWarning: Model failed to serialize as JSON. Ignoring... Invalid format specifier
  warnings.warn(f"Model failed to serialize as JSON. Ignoring... {exc}")

要从 Jupyter notebook 加载 TensorBoard,您可以使用 %tensorboard 魔法命令

%tensorboard --logdir logs

TensorBoard 监控指标并检查训练曲线。

Tensorboard training graph

TensorBoard 还允许您探索模型中使用的计算图。

Tensorboard graph exploration

在调试过程中,能够深入了解您的模型非常有价值。


结论

使用 tensorflow_numpy API 将现有的 NumPy 代码移植到 Keras 模型很容易!通过与 Keras 集成,您可以使用现有的 Keras 回调、指标和优化器,轻松分发您的训练并使用 Tensorboard。

将更复杂的模型(例如 ResNet)迁移到 TensorFlow NumPy API 将是一个很好的后续学习练习。

在线提供了一些开源 NumPy ResNet 实现。