当今计算机科学和机器学习 (ML) 的许多激动人心的应用都处理跨单个大型坐标系的多维数据集，例如，根据空间网格上的大气测量值进行天气建模，或根据 2d 或 3d 扫描中的多通道图像强度值进行医学成像预测。在这些设置中，即使是单个数据集也可能需要 TB 或 PB 的数据存储。处理此类数据集也具有挑战性，因为用户可能会以不规则的间隔和不同的规模读取和写入数据，并且通常对使用多台并行工作的机器进行分析感兴趣。
今天，我们介绍TensorStore，这是一个开源 C++ 和 Python 软件库，专为存储和处理n维数据而设计：
提供统一的 API 用于读写多种数组格式，包括zarr和N5。
原生支持多种存储系统，包括Google Cloud Storage、本地和网络文件系统、HTTP 服务器和内存存储。
支持读/写回缓存和事务，具有强大的原子性、隔离性、一致性和持久性(ACID) 保证。
通过乐观并发支持从多个进程和机器进行安全、高效的访问。
提供异步 API，即使对高延迟的远程存储也能实现高吞吐量访问。
提供先进的、完全可组合的索引操作和虚拟视图。
TensorStore 已被用于解决科学计算中的关键工程挑战（例如，神经科学中大型数据集的管理和处理，如千万亿级 3D 电子显微镜数据和神经元活动的“4D”视频）。TensorStore 还用于创建大规模机器学习模型，如PaLM，以解决分布式训练期间管理模型参数（检查点）的问题。
熟悉的数据访问和操作 API
TensorStore 提供了一个简单的 Python API，用于加载和操作大型数组数据。在下面的示例中，我们创建一个 TensorStore 对象，该对象表示苍蝇大脑的 56 万亿体素 3D 图像，并以NumPy数组的形式访问 100x100 的小块数据：
>>> 将 TensorStore 导入为 ts
>>> 将 numpy 导入为 np
# 创建一个 TensorStore 对象来处理苍蝇脑数据。
>>> 数据集 = ts.open({
... '司机'：
...'neuroglancer_precomputed'，
...'kvstore'：
... 'gs://neuroglancer-janelia-flyem-hemibrain/' +
...'v1.1/分段/'，
... }）。结果（）
# 创建 3-d 视图（删除单例“通道”维度）：
>>> dataset_3d = 数据集[ts.d['channel'][0]]
>>> 数据集_3d.域
{ “x”：[0，34432)，“y”：[0，39552)，“z”：[0，41408) }
# 将 100x100x1 的数据切片转换为 numpy ndarray
>>>切片=np.array（dataset_3d[15000:15100, 15000:15100, 20000]）
至关重要的是，在请求特定的 100x100 切片之前，不会访问或存储任何实际数据；因此，可以加载和操作任意大小的底层数据集，而不必将整个数据集存储在内存中，使用与标准 NumPy 操作大致相同的索引和操作语法。TensorStore 还为高级索引功能提供了广泛的支持，包括转换、对齐、广播和虚拟视图（数据类型转换、下采样、延迟动态生成的数组）。
以下示例演示了如何使用 TensorStore 创建 zarr 数组，以及其异步 API 如何实现更高的吞吐量：
>>> 将 TensorStore 导入为 ts
>>> 将 numpy 导入为 np
>>> # 在本地文件系统上创建一个 zarr 数组
>>> 数据集 = ts.open({
...'司机'：'zarr'，
...'kvstore'：'file:///tmp/my_dataset/'，
... },
... dtype=ts.uint32，
...chunk_layout=ts.ChunkLayout(chunk_shape=[256, 256, 1]),
...创建=True，
...形状=[5000，6000，7000]）.result（）
>>> # 创建两个带有要写入的示例数据的 numpy 数组。
>>> a = np.arange(100*200*300, dtype=np.uint32).reshape((100, 200, 300))
>>> b = np.arange(200*300*400, dtype=np.uint32).reshape((200, 300, 400))
>>> # 启动两个异步写入，同时执行。
>>> future_a = 数据集[1000:1100, 2000:2200, 3000:3300].写入（a）
>>> future_b = 数据集[3000:3200, 4000:4300, 5000:5400].写入（b）
>>> # 等待异步写入完成
>>> Future_a.result()
>>> future_b.result()
安全且高效的扩展
处理和分析大型数值数据集需要大量计算资源。这通常是通过分布在多台机器上的大量 CPU 或加速器核心的并行化来实现的。因此，TensorStore 的一个基本目标是实现单个数据集的并行处理，既安全（即避免由并行访问模式引起的损坏或不一致），又具有高性能（即在计算过程中，对 TensorStore 的读写不会成为瓶颈）。事实上，在 Google 数据中心的测试中，我们发现随着 CPU 数量的增加，读写性能几乎呈线性增长：
使用 Google 数据中心中可变数量的单核计算任务并发访问驻留在 Google Cloud Storage (GCS) 上的zarr 格式的 TensorStore 数据集的读取和写入性能。读取和写入性能都与计算任务的数量几乎呈线性关系。
性能的实现方式是：使用 C++ 实现核心操作，广泛使用多线程执行编码/解码和网络 I/O 等操作，并通过分块将大型数据集划分为更小的单元，从而高效地读取和写入整个数据集的子集。TensorStore 还提供可配置的内存缓存（可减少频繁访问数据的较慢存储系统交互）和异步 API，使读取或写入操作可以在程序完成其他工作时在后台继续进行。
当多台机器访问同一数据集时，通过使用乐观并发来实现并行操作的安全性，这可以保持与各种底层存储层（包括云存储平台，例如GCS以及本地文件系统）的兼容性，而不会显着影响性能。TensorStore 还为在单个运行时内执行的所有单个操作提供强大的 ACID 保证。
为了使 TensorStore 的分布式计算与许多现有的数据处理工作流程兼容，我们还将 TensorStore 与并行计算库（如Apache Beam（示例代码）和Dask（示例代码））集成在一起。
用例：语言模型
机器学习领域一个令人兴奋的最新进展是更先进的语言模型（如PaLM）的出现。这些神经网络包含数千亿个参数，在自然语言理解和生成方面表现出一些令人惊讶的能力。这些模型还突破了计算基础设施的极限；特别是，训练像 PaLM 这样的语言模型需要数千个 TPU 并行工作。
在此训练过程中出现的一个挑战是如何高效地读取和写入模型参数。训练分布在许多独立的机器上，但必须定期将参数保存到永久存储系统上的单个对象（“检查点”）中，而不会减慢整个训练过程。各个训练作业还必须能够只读取它们关注的特定参数集，以避免加载整个模型参数集（可能高达数百 GB）所需的开销。
TensorStore 已用于解决这些挑战。它已用于管理与使用JAX （代码示例）训练的大规模（“多足”）模型相关的检查点，并已与T5X （代码示例）和Pathways等框架集成。模型并行性用于在数百个 TPU 上对整个参数集进行分区，这些参数集可能占用超过 1TB 的内存。使用 TensorStore 将检查点存储为 zarr 格式，并选择块结构以允许并行独立读取和写入每个 TPU 的分区。
保存检查点时，每个模型参数都使用 TensorStore 以 zarr 格式写入，使用块网格进一步细分用于在 TPU 上对参数进行分区的网格。主机并行写入分配给连接到该主机的 TPU 的每个分区的 zarr 块。使用 TensorStore 的异步 API，即使数据仍在写入持久存储，训练仍会继续进行。从检查点恢复时，每个主机仅读取组成分配给该主机的分区的块。
用例：3D 脑图
突触分辨率连接组学 领域旨在以单个突触连接的详细级别绘制动物和人类大脑的布线。这需要在高达毫米或更大的视野范围内以极高的分辨率（纳米）对大脑进行成像，从而产生可以跨越拍字节大小的数据集。未来，随着科学家考虑绘制整个小鼠或灵长类动物的大脑，这些数据集可能会扩展到艾字节。然而，即使是当前的数据集也面临着与存储、操作和处理相关的重大挑战；特别是，即使是单个大脑样本也可能需要数百万千兆字节，其坐标系（像素空间）在每个维度上都有数十万像素。
我们已经使用 TensorStore 解决了与大规模连接组数据集相关的计算挑战。具体来说，TensorStore 管理着一些最大、访问最广泛的连接组数据集，并以 Google Cloud Storage 作为底层对象存储系统。例如，它已应用于人类皮层“h01”数据集，该数据集是人类脑组织的 3D 纳米分辨率图像。原始成像数据为 1.4 PB（大约 500,000 * 350,000 * 5,000 像素大），并且还与驻留在相同坐标系中的其他内容（如 3D 分割和注释）相关联。原始数据被细分为 128x128x16 像素大的单个块，并以“ Neuroglancer 预计算”格式存储，该格式针对基于 Web 的交互式查看进行了优化，并且可以轻松地从 TensorStore 中进行操作。
苍蝇大脑重建，使用 TensorStore可以轻松访问和操作其底层数据。
致谢
感谢Google Research 的Tim Blakely、Viren Jain、Yash Katariya、Jan-Matthis Luckmann、Michał Januszewski、Peter Li、Adam Roberts、Brain Williams和Hector Yee ，以及广大科学界的Davis Bennet、Stuart Berg、Eric Perlman、Stephen Plaza和Juan Nunez-Iglesias在设计、早期测试和调试方面提供的宝贵反馈。
评论