导入
配置 TensorBoard
学习如何在 JAX 和 Flax 中使用 TensorBoard 来训练模型。
跟踪机器学习实验有助于轻松理解并可视化模型的性能，同时也能帮助发现网络中的潜在问题。例如，通过观察训练和验证曲线，我们可以快速识别过拟合现象。这些图表可以使用我们熟悉的绘图工具包（如 Matplotlib）绘制，但也可以采用更先进的工具，比如 TensorBoard。
...
...
TensorBoard 是一个开源库，提供用于机器学习实验跟踪的一系列工具。我们可以使用 TensorBoard 实现以下功能：
- 跟踪和可视化模型评估指标，例如准确率。
- 记录图像数据。
- 可视化超参数调优过程。
- 投影嵌入向量，例如自然语言处理中的词嵌入（word embedding）。
- 可视化模型权重和偏置项的直方图。
- 绘制模型结构图。
- 分析网络性能。
TensorBoard 支持多种主流机器学习库，包括 XGBoost、JAX、Flax 和 PyTorch。
本节课将重点介绍如何在使用 JAX 和 Flax 构建神经网络时集成 TensorBoard。
在 Jupyter Notebook 和 Google Colaboratory 中使用 TensorBoard
首先，在运行环境中加载 TensorBoard 扩展，通常是在 Google Colaboratory 或本地 Jupyter Notebook 中执行：
1%load_ext tensorboard
接着，指定用于存储日志信息的目录：
1log_folder = "runs"
启动 TensorBoard
在 Notebook 环境中，可通过 tensorboard 魔法命令启动 TensorBoard，并指定日志目录 logdir：
1%tensorboard --logdir={log_folder}
TensorBoard 仪表板
TensorBoard 提供多个仪表板，用于展示不同类型的信息：
- Scalars 仪表板：用于跟踪数值型指标，例如每个训练周期的损失和准确率。也可用来监控其他标量值，如训练速度或学习率的变化。
...
TensorBoard 的 Scalars 仪表板
- Graphs 仪表板：用于展示模型的可视化结构。例如，可用于检查网络的整体架构。
...
TensorBoard 的 Graphs 仪表板
- Distributions（分布）和 Histograms（直方图）仪表板：显示张量随时间变化的分布情况，常用于监控网络中权重和偏置项的分布状态。
...
TensorBoard 的 Distributions 仪表板
...
TensorBoard 的 Histograms 仪表板
- Images 仪表板：展示已记录到 TensorBoard 中的图像数据。
...
TensorBoard 的 Images 仪表板
- HParams 仪表板：用于可视化超参数优化过程，帮助识别最优的网络参数组合。
...
TensorBoard 的 HParams 仪表板
- Embedding Projector（嵌入投影器）：用于可视化低维空间中的嵌入表示，例如文本嵌入。
...
TensorBoard 的 Embedding Projector
- What-If Tool 仪表板：帮助分析模型性能，支持对不同数据子集进行公平性评估。
...
TensorBoard 的 What-If Tool 仪表板
- TensorFlow Profiler：监控模型训练过程，展示 CPU 和 GPU 上的运行事件，并基于采集的数据提供性能优化建议。此外，还可用于调试输入流水线中的性能瓶颈。
...
TensorBoard 的 Profile 仪表板