谷歌官方高性能大规模高维数据处理库TensorStore发布！

标签：#Tensorflow##TensorStore##好物分享##高维数据# 时间：2022/09/24 10:15:50 作者：小木

许多令人振奋的当代计算机科学和机器学习（ML）的应用都在处理跨越单一大坐标系的多维数据集，例如，从空间网格上的大气测量结果进行天气建模，或从2D或3D扫描的多通道图像强度值进行医学成像预测。在这些情况下，即使是一个单一的数据集也可能需要TB级或PB级的数据存储。这样的数据集在工作中也具有挑战性，因为用户可能以不规则的时间间隔和不同的规模来读写数据，而且往往对使用众多机器并行工作来进行分析感兴趣。

今天Google发布了TensorStore，这是一个开源的C++和Python软件库，设计用于存储和操作大规模n维数据，它包含如下特点：

为读写多种阵列格式提供统一的API，包括zarr和N5。
天然支持多种存储系统，包括谷歌云存储、本地和网络文件系统、HTTP服务器和内存存储。
支持读/写缓存和事务，具有强大的原子性、隔离性、一致性和耐久性（ACID）保证。
通过乐观的并发性，支持来自多个进程和机器的安全、高效的访问。
提供一个异步的API，以实现高吞吐量的访问，甚至是对高延迟的远程存储。
提供先进的、完全可组合的索引操作和虚拟视图。

TensorStore已经被用来解决科学计算中的关键工程挑战（例如，管理和处理神经科学中的大型数据集，如石油级的三维电子显微镜数据和神经元活动的 “4d “视频）。TensorStore还被用于创建大规模的机器学习模型，如PaLM，解决了分布式训练期间管理模型参数（检查点）的问题。

熟悉的数据访问和操作的API

TensorStore提供了一个简单的Python API，用于加载和操作大型阵列数据。在下面的例子中，我们创建了一个TensorStore对象，它代表了一个56万亿体素的苍蝇大脑三维图像，并以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还提供了对高级索引功能的广泛支持，包括转换、对齐、广播和虚拟视图（数据类型转换、下采样、懒散地即时生成数组）。

安全和性能的扩展

处理和分析大型数值数据集需要大量的计算资源。这通常是通过分布在许多机器上的许多CPU或加速器核心的并行化实现的。因此，TensorStore的一个基本目标是使单个数据集的并行处理既安全（即避免并行访问模式引起的损坏或不一致）又高性能（即在计算过程中，对TensorStore的读写不会成为瓶颈）。事实上，在谷歌数据中心内的测试中，我们发现随着CPU数量的增加，读写性能几乎呈线性扩展。

此外，Google还展示了TensorStore在其它方面的应用。

TensorStore在大规模语言模型上的使用

ML的一个令人兴奋的最新发展是出现了更先进的语言模型，如PaLM。这些神经网络包含数以千亿计的参数，在自然语言理解和生成方面表现出一些令人惊讶的能力。这些模型也推动了计算基础设施的极限；特别是，训练像PaLM这样的语言模型需要成千上万的TPU并行工作。

在这个训练过程中出现的一个挑战是有效地读取和写入模型参数。训练分布在许多独立的机器上，但参数必须定期保存到永久存储系统上的单一对象（”检查点”），而不拖累整个训练过程。单独的训练工作也必须能够只读取他们所关注的特定参数集，以避免加载整个模型参数集（可能是数百GB）的开销。

TensorStore已经被用来解决这些挑战。它已被应用于管理与用JAX训练的大规模（”multipod”）模型相关的检查点（代码示例），并已与T5X（代码示例）和Pathways等框架集成。模型的并行性被用来划分全套参数，这些参数可能占据超过一兆字节的内存，在数百个TPU上。检查点使用TensorStore以zarr格式存储，选择的块状结构允许每个TPU的分区以并行的方式独立读取和写入。

TensorStore在3维大脑图谱上的应用

我们已经使用TensorStore来解决与大规模连接组数据集相关的计算挑战。具体来说，TensorStore已经管理了一些最大和最广泛访问的连接体数据集，以谷歌云存储作为基础对象存储系统。例如，它已被应用于人类皮质 “h01 “数据集，这是人类大脑组织的三维纳米分辨率图像。原始成像数据是1.4PB（大约500,000 350,000 5,000像素大），并与驻扎在同一坐标系中的三维分割和注释等附加内容进一步关联。原始数据被细分为128x128x16像素大小的独立块，并以 “Neuroglancer预计算 “的格式存储，该格式为基于网络的交互式查看进行了优化，并可从TensorStore轻松操作。

大家可以自行下载去测试使用TensorStore：https://github.com/google/tensorstore