使用 TensorFlow NumPy 编写 Keras 模型
作者： lukewood
创建日期 2021/08/28
上次修改日期 2021/08/28
描述：概述如何使用 TensorFlow NumPy API 编写 Keras 模型。
简介
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 还允许您探索模型中使用的计算图。
在调试期间，能够洞察您的模型可能非常有价值。
结论
使用 tensorflow_numpy API 将现有的 NumPy 代码移植到 Keras 模型非常容易！通过与 Keras 集成，您可以使用现有的 Keras 回调、指标和优化器，轻松分发您的训练并使用 Tensorboard。
将更复杂的模型（例如 ResNet）迁移到 TensorFlow NumPy API 将是一个很好的后续学习练习。
网上有几个开源的 NumPy ResNet 实现可用。