JupyterLab TensorBoard Pro:功能完备的 TensorBoard 插件

High-Flyer    August 01, 2022

对于深度学习开发者和研究者们来说,TensorBoard 是帮助理解深度学习运行机理,训练可视化和调试的重要工具。

最近,幻方 AI 开发并集成了 JupyterLab TensorBoard 插件(JupyterLab TensorBoard Pro),帮助大家在使用幻方萤火超算集群训练模型的同时,也可以方便快捷的使用 Tensorboard 工具进行深度学习训练的调试。

该插件更新修复了社区里原插件的若干问题,并对性能进行了提升。我们对该插件进行了开源,方便用户在没有幻方萤火集群的 JupyterLab 环境的情况下,仍然可以下载和使用。下载安装方式如下

pip install jupyterlab-tensorboard-pro

开源地址https://github.com/HFAiLab/jupyterlab_tensorboard_pro

研发背景

目前深度学习社区里面已经有了 jupyterlab_tensorboard(前端插件)和 jupyter_tensorboard(对应的后端插件)两个开源插件,然而这两个仓库已经很久没有更新迭代,对于一些外部提交的PR也没有及时合入(实际上,当我们发布 jupyterlab_tensorboard_pro 之后,jupyterlab_tensorboard 合入了一些社区 PR,但 jupyter_tensorboard 仍然存在无法配合较新的 tensorboard 版本使用的问题),基于此我们判断该项目作者已经不在积极地维护对应的仓库。

同时,现有社区的 TensorBoard 插件存在一定的体验问题:

  1. 需要同时安装两个 Python 包,并且需要版本对应,增加了使用和维护的成本;
  2. 社区版本插件启动 TensorBoard 后端会持续 reload(从磁盘读取文件并更新),这会对 JupyterLab 的体验产生副作用,如果使用了分布式文件系统或者公用集群,甚至还会影响到其他用户。目前社区版本插件无法通过配置来禁用。
  3. 功能体验不够完善,例如点击 TensorBoard 按钮之后会阻塞请求导致 JupyterLab 没有响应,并且没有任何提示,同时也对已经创建的 TensorBoard 实例的可视化管理功能也比较缺乏。

基于以上原因,幻方AI基于现有的社区版本二次开发了 TensorBoard 插件,以解决上述问题。在内部使用一段时间后,插件运行稳定并持续获得正向反馈,因此我们将插件开源以帮助更多的开发者和研究者们。

在萤火 JupyterLab 中使用

创建实例

1. 从 launcher 面板创建

您可以登录进入萤火 JupyterLab,从 Launcher 面板点击 TensorBoard 图标。首次点击会进入到一个初始化面板,您可以在该面板内创建 TensorBoard 实例。非首次进入则会直接进入到第一个活跃的 TensorBoard 实例。

tensorboard step1

2. 通过快捷命令创建

您也可以打开 JupyterLab 快捷指令面板(ctrl + shift + c 唤起),输入 Open TensorBoard

tensorboard step2

设置参数

这里,在初始化面板中,提供了两个参数可以进行设置:

  • Log Dir:日志目录地址,默认是点击 TensorBoard 时当前侧边栏的目录。您也可以手动填写对应目录。这里建议目录尽可能的细化,因为目录内容少会提高初始化的速度。
  • Reload Interval:TensorBoard 多久对对应目录进行一次重新扫描,默认是 120s。如果您不需要扫描建议及时关闭,日常使用手动 Reload 即可。这里要注意的是:设置 Reload Interval 之后,TensorBoard 后端持续扫描目录会对 JupyterLab 的稳定性和文件系统都产生一定的影响。

选择好参数点击 Create TensorBoard,会同步创建 TensorBoard 实例,这个时候 JupyterLab 后端是阻塞的,请等待实例创建好之后再进行其他操作。

tensorboard step3

管理实例

创建好实例后,您可以对 TensorBoard 的实例进行管理。目前提供了以下几个功能:

  • 刷新和列表切换:可以切换成其他的实例的 TensorBoard 后端,这个时候不会销毁实例。
  • 独立页面中打开:可以以独立网页 Tab 的形式打开 TensorBoard。
  • Reload:即重新初始化 TensorBoard 后端,当文件内容有更新时,可以通过此功能载入新的内容(注:TensorBoard 内部的刷新,不会造成 Reload)。
  • Destroy:销毁实例,会连同前端面板一起关掉。
  • Duplicate:重新打开一个完全一样的前端面板,此操作会复用 TensorBoard 后端。
  • New:额外新建一个 TensorBoard 后端,注意事项可以参考上文。

另外,对于我们创建的 TensorBoard 实例,可以在 JupyterLab 的 Kernel 管理面板一同管理,提供跳转至对应实例和删除等功能。

tensorboard step4

在本地安装使用

JupyterLab-TensorBoard-Pro 支持的最低 Python 版本为 Python 3.6,不过由于幻方 AI 套件建议的 Python 版本为 3.8,因此建议您在 Python 3.8 的环境中使用该插件,从而方便安装其他幻方 AI 套件。

为了兼容多版本的 TensorBoard 和 JupyterLab,插件本身未指定强制依赖版本,需要提前确认本地环境已经安装如下第三方库:

  1. jupyterlab(V3+)
  2. tensorflow(建议 V2+)
  3. tensorboard(建议 V2+)

如果是一个全新的 Python 环境,可以直接 pip install jupyterlab tensorflow tensorboard 安装当前默认版本即可。

通过 pip install jupyterlab-tensorboard-pro 下载安装之后,会默认安装到当前 Python 环境的 JupyterLab 插件列表中,直接启动 JupyterLab 即可使用,使用方式和注意事项与在萤火 JupyterLab 中使用相同。

总结

JupyterLab-TensorBoard-Pro 旨在提供高性能的 TensorBoard 插件服务,帮助您更方便地在萤火集群中调优深度学习训练。我们欢迎大家来试用,并提出反馈


本文作者: High-Flyer


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。