计算机科学和机器学习 (ML) 的许多令人兴奋的当代应用程序操纵跨越单个大坐标系的多维数据集，例如，通过空间网格上的大气测量进行天气建模或从 2d 中的多通道图像强度值进行医学成像预测或 3D 扫描。在这些设置中，即使是单个数据集也可能需要 TB 或 PB 的数据存储。使用此类数据集也具有挑战性，因为用户可能会以不规则的间隔和不同的规模读取和写入数据，并且通常对使用大量并行工作的机器执行分析感兴趣。
今天我们将介绍TensorStore，这是一个开源 C++ 和 Python 软件库，专为存储和操作n维数据而设计，它：
- 提供统一的 API 用于读写多种数组格式，包括zarr和N5。
- 本机支持多种存储系统，包括Google Cloud Storage、本地和网络文件系统、HTTP 服务器和内存存储。
- 支持读/写缓存和事务，具有强大的原子性、隔离性、一致性和持久性(ACID) 保证。
- 通过并发支持从多个进程和机器进行安全、高效的访问。
- 提供异步 API 以实现对高延迟远程存储的高吞吐量访问。
- 提供高级、完全可组合的索引操作和虚拟视图。
TensorStore 已被用于解决科学计算中的关键工程挑战（例如，管理和处理神经科学中的大型数据集，例如peta 级 3d 电子显微镜数据和神经元活动的“4d”视频）。通过解决分布式训练期间管理模型参数（检查点）的问题， TensorStore 还被用于创建大型机器学习模型，例如PaLM 。
数据访问和操作使用相似的API
TensorStore 提供了一个简单的 Python API，用于加载和操作大型数组数据。在下面的示例中，我们创建了一个 TensorStore 对象，该对象代表一个 56 万亿体素的苍蝇大脑 3d 图像，并以NumPy数组 的形式访问一小块 100x100 的数据块：
>>>?import?tensorstore?as?ts
>>>?import?numpy?as?np
#?Create?a?TensorStore?object?to?work?with?fly?brain?data.
>>>?dataset?=?ts.open({
...?????'driver':
...?????????'neuroglancer_precomputed',
...?????'kvstore':
...?????????'gs://neuroglancer-janelia-flyem-hemibrain/'?+?
...?????????'v1.1/segmentation/',
...?}).result()
#?Create?a?3-d?view?(remove?singleton?'channel'?dimension):
>>>?dataset_3d?=?dataset[ts.d['channel'][0]]
>>>?dataset_3d.domain
{?"x":?[0,?34432),?"y":?[0,?39552),?"z":?[0,?41408)?}
#?Convert?a?100x100x1?slice?of?the?data?to?a?numpy?ndarray
>>>?slice?=?np.array(dataset_3d[15000:15100,?15000:15100,?20000])
至关重要的是，在请求特定的 100x100 切片之前，不会访问或存储在内存中的实际数据；因此，可以加载和操作任意大的基础数据集，而无需将整个数据集存储在内存中，使用与标准 NumPy 操作基本相同的索引和操作语法。TensorStore 还为高级索引功能提供广泛的支持，包括转换、对齐、广播和虚拟视图（数据类型转换、下采样、延迟动态生成的数组）。
以下示例演示了如何使用 TensorStore 创建一个 zarr 数组，以及它的异步 API 如何实现更高的吞吐量：
>>>?import?tensorstore?as?ts
>>>?import?numpy?as?np
>>>?#?Create?a?zarr?array?on?the?local?filesystem
>>>?dataset?=?ts.open({
...?????'driver':?'zarr',
...?????'kvstore':?'file:///tmp/my_dataset/',
...?},
...?dtype=ts.uint32,
...?chunk_layout=ts.ChunkLayout(chunk_shape=[256,?256,?1]),
...?create=True,
...?shape=[5000,?6000,?7000]).result()
>>>?#?Create?two?numpy?arrays?with?example?data?to?write.
>>>?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))
>>>?#?Initiate?two?asynchronous?writes,?to?be?performed?concurrently.
>>>?future_a?=?dataset[1000:1100,?2000:2200,?3000:3300].write(a)
>>>?future_b?=?dataset[3000:3200,?4000:4300,?5000:5400].write(b)
>>>?#?Wait?for?the?asynchronous?writes?to?complete
>>>?future_a.result()
>>>?future_b.result()
安全和高性能的平衡
处理和分析大型数值数据集需要大量的计算资源。这通常是通过分布在许多机器上的众多 CPU 或加速器内核的并行化来实现的。因此，TensorStore 的一个基本目标是实现对单个数据集的并行处理，这些数据集既安全（即避免由并行访问模式引起的损坏或不一致）又具有高性能（即，对 TensorStore 的读写不是计算过程中的瓶颈） . 事实上，在 Google 数据中心内的一项测试中，我们发现随着 CPU 数量的增加，读写性能几乎呈线性增长：
驻留在谷歌云存储 (GCS) 上的zarr 格式的 TensorStore 数据集的读写性能，使用谷歌数据中心中可变数量的单核计算任务同时访问。读写性能几乎与计算任务的数量呈线性关系
性能是通过在 C++ 中实现核心操作、在编码/解码和网络 I/O 等操作中广泛使用多线程以及通过分块将大型数据集划分为更小的单元来实现高效读取和写入整个数据集的子集来实现的。TensorStore 还提供可配置的内存缓存（可减少频繁访问的数据的较慢存储系统交互）和异步 API，允许在程序完成其他工作时在后台继续进行读取或写入操作。
当多台机器访问同一个数据集时，并行操作的安全性是通过使用乐观并发来实现的，它保持与各种底层存储层（包括云存储平台，如GCS以及本地文件系统）的兼容性，而不会显着影响性能。TensorStore 还为在单个运行时中执行的所有单个操作提供强大的 ACID 保证。
为了使使用 TensorStore 的分布式计算与许多现有的数据处理工作流程兼容，我们还将 TensorStore 与Apache Beam（示例代码）和Dask（示例代码）等并行计算库集成。
案例: 语言模型
ML 最近一个令人兴奋的发展是出现了更高级的语言模型，例如PaLM。这些神经网络包含数千亿个参数，并在自然语言理解和生成方面表现出一些令人惊讶的能力。这些模型也突破了计算基础设施的极限；特别是，训练诸如 PaLM 之类的语言模型需要数千个 TPU 并行工作。
在此训练过程中出现的一个挑战是有效地读取和写入模型参数。训练分布在许多不同的机器上，但参数必须定期保存到永久存储系统上的单个对象（“检查点”），而不会减慢整个训练过程。单个训练作业还必须能够仅读取他们关注的特定参数集，以避免加载整个模型参数集（可能是数百 GB）所需的开销。
TensorStore 已经被用来应对这些挑战。它已被应用于管理与使用JAX （代码示例）训练的大型（“多脚架”）模型相关的检查点，并已与T5X （代码示例）和Pathways等框架集成。模型并行性用于划分完整的参数集，这些参数可以占用超过 TB 的内存，超过数百个 TPU。检查点使用 TensorStore 以 zarr 格式存储，并选择了一个块结构以允许每个 TPU 的分区并行独立地读取和写入。
保存检查点时，每个模型参数都使用 zarr 格式的 TensorStore 写入，使用块网格进一步细分用于在 TPU 上划分参数的网格。主机并行写入分配给连接到该主机的 TPU 的每个分区的 zarr 块。使用 TensorStore 的异步 API，即使数据仍在写入持久存储，训练也会继续进行。从检查点恢复时，每个主机只读取组成分配给该主机的分区的块
案例: 大脑3D映射
突触分辨率连接组学领域旨在在单个突触连接的详细级别绘制动物和人类大脑的布线。这需要在高达毫米或更大的视野范围内以极高的分辨率（纳米）对大脑进行成像，从而产生可以跨越 PB 大小的数据集。未来，随着科学家考虑绘制整个鼠标，这些数据集可能会扩展到 EB 级或灵长类动物的大脑。然而，即使是当前的数据集也带来了与存储、操作和处理相关的重大挑战。特别是，即使是单个大脑样本也可能需要数百万千兆字节，每个维度的坐标系统（像素空间）都有数十万像素。
我们使用 TensorStore 来解决与大规模连接组数据集相关的计算挑战。具体来说，TensorStore 管理了一些最大和最广泛访问的连接组数据集，其中谷歌云存储作为底层对象存储系统。例如，它已应用于人类皮层“h01”数据集，这是人类脑组织的 3d 纳米分辨率图像。原始成像数据为 1.4 PB（大约 500,000 * 350,000 * 5,000 像素大，并进一步与位于同一坐标系中的 3d 分割和注释等其他内容相关联。原始数据被细分为 128x128x16 像素大的单个块，以“ Neuroglancer precomputed ”格式存储，该格式针对基于 Web 的交互式查看，可以从 TensorStore 轻松操作。
一种苍蝇大脑重建，可以使用 TensorStore轻松访问和操作其基础数据
安装方法
要开始使用 TensorStore Python API，您可以使用以下命令安装tensorstore PyPI 包：
pip?install?tensorstore