当前位置: 首页 > 科技观察

存储和操作n维数据的难题,谷歌用一个开源软件库解决了

时间:2023-03-17 16:15:50 科技观察

存储和操作n维数据的难题,Google使用开源软件库解决了TB或PB级数据。另一方面,使用此类数据集也具有挑战性,因为用户可能会以不规则的间隔和不同的规模读写数据,通常会执行大量并行工作。为了解决上述问题,Google开发了TensorStore,这是一个开源的C++和Python软件库,专为存储和操作n维数据而设计。GoogleAI负责人JeffDean也在推特上表示,TensorStore现已正式开源。TensorStore的主要功能包括:提供统一的API读写多种数组格式,包括zarr和N5;原生支持多种存储系统,包括谷歌云存储、本地和网络文件系统、HTTP服务器和内存存储;支持读/写缓存和事务,具有强原子性、隔离性、一致性和持久性(ACID)特性;支持多进程多机安全高效并发访问;提供异步API,实现高延迟远程高吞吐存储访问;提供高级的、完全可组合的索引操作和虚拟视图。TensorStore已被用于解决科学计算中的工程挑战,也被用于创建大规模机器学习模型,例如在PaLM的分布式训练期间管理模型参数(检查点)。GitHub地址:https://github.com/google/tensorstore用于数据访问和操作的APITensorStore提供了一个简单的PythonAPI,用于加载和操作大型数组数据。例如,下面的代码创建了一个TensorStore对象,它表示一个56万亿体素的苍蝇大脑3D图像,并允许访问NumPy数组中的100x100图像补丁数据:在补丁之前,不会访问内存中的实际数据,因此任意无需将整个数据集存储在内存中即可加载和操作大型基础数据集。TensorStore使用与标准NumPy基本相同的索引和操作语法。此外,TensorStore还为高级索引功能提供广泛支持,包括对齐、虚拟视图等。以下代码演示了如何使用TensorStore创建一个zarr数组,以及TensorStore的异步API如何实现更高的吞吐量:安全性和性能扩展众所周知,分析和处理大数据集需要大量的计算资源,通常需要分布在机器上CPU或加速器内核的多个并行化上。因此,TensorStore的一个基本目标就是实现并行处理,达到安全和高性能的兼顾。事实上,在谷歌数据中心的测试中,他们发现随着CPU数量的增加,TensorStore的读写性能几乎呈线性增长:zarr格式数据集在谷歌云存储(GCS)上的读写性能,读写性能随着计算任务的数量几乎呈线性增长。TensorStore还提供可配置的内存缓存和异步API,以允许在程序执行其他工作时在后台继续读取和写入。为了使TensorStore的分布式计算与数据处理工作流兼容,谷歌还将TensorStore与ApacheBeam等并行计算库集成在一起。例子展示例1语言模型:最近,机器学习领域出现了一些高级语言模型,如PaLM。这些模型包含数千亿个参数,并在自然语言理解和生成方面展现出惊人的能力。然而,这些模型对计算设施提出了挑战,特别是训练像PaLM这样的语言模型需要数千个TPU并行工作。其中,有效读写模型参数是训练过程面临的一个问题:比如训练分布在不同的机器上,但必须每隔一定时间将参数保存到checkpoint;另一个例子是单次训练必须只读取特定的参数集,以避免加载整个模型参数集(可能是数百GB)所需的开销。TensorStore可以解决以上问题。它已被用于管理与大型(多脚架)模型相关的检查点,并已与T5X和Pathways等框架集成。TensorStore将Checkpoint转换为zarr格式存储,并选择block结构,让每个TPU的partition可以独立并行读写。保存checkpoint时,将参数写入zarr格式,进一步划分blockgrid,用于在TPU上划分参数grid。主机为分配给该主机的TPU的每个分区并行写入zarr块。使用TensorStore的异步API,即使数据仍在写入持久存储,训练也会继续。从检查点恢复时,每个主机只读取分配给该主机的分区块。示例2大脑3D映射:突触解析连接组学旨在在个体突触连接水平上绘制动物和人类大脑的线路图。实现这一目标需要在毫米或更大的视野范围内以极高的分辨率(纳米级)对大脑进行成像,从而产生PB级的数据。然而,即使是现在,数据集也面临存储、处理等问题,即使是单个大脑样本也可能需要数百万GB的空间。谷歌使用TensorStore来解决与大规模连接组学数据集相关的计算挑战。具体来说,TensorStore已经开始管理一些连接组学数据集,并使用谷歌云存储作为底层对象存储系统。目前,TensorStore已用于人类大脑皮层数据集H01,原始成像数据为1.4PB(约500000*350000*5000像素)。然后将原始数据细分为128x128x16像素的单独块,以“Neuroglancer预计算”格式存储,可以通过TensorStore轻松操作。TensorStore可用于轻松访问和操作底层数据(飞脑重建)。想要入门的朋友可以使用以下方法安装TensorStorePyPI包:pipinstalltensorstore