計(jì)算機(jī)科學(xué)和機(jī)器學(xué)習(xí) (ML) 的許多應(yīng)用都需要處理跨坐標(biāo)系的多維數(shù)據(jù)集，并且單個(gè)數(shù)據(jù)集可能也需要存儲(chǔ) TB 或 PB 的數(shù)據(jù)。另一方面，使用此類數(shù)據(jù)集也具有挑戰(zhàn)性，因?yàn)橛脩艨赡軙?huì)以不規(guī)則的間隔和不同的規(guī)模讀取和寫入數(shù)據(jù)，通常還會(huì)執(zhí)行大量的并行工作。

為了解決上述問(wèn)題，谷歌開(kāi)發(fā)了一個(gè)開(kāi)源的 C++ 和 Python 軟件庫(kù) TensorStore，專為存儲(chǔ)和操作 n 維數(shù)據(jù)而設(shè)計(jì)。谷歌 AI 負(fù)責(zé)人 Jeff Dean 也在推特上發(fā)文表示 TensorStore 現(xiàn)已正式開(kāi)源。

TensorStore 的主要功能包括：

• 提供統(tǒng)一的 API 用于讀寫多種數(shù)組格式，包括 zarr 和 N5；
• 原生支持多種存儲(chǔ)系統(tǒng)，包括谷歌云存儲(chǔ)、本地和網(wǎng)絡(luò)文件系統(tǒng)、HTTP 服務(wù)器和內(nèi)存存儲(chǔ)；
• 支持讀 / 寫緩存和事務(wù)，具有很強(qiáng)的原子性、隔離性、一致性和持久性（ACID）特性；
• 支持從多個(gè)進(jìn)程和機(jī)器進(jìn)行安全、高效的并發(fā)訪問(wèn)；
• 提供異步 API 以實(shí)現(xiàn)對(duì)高延遲遠(yuǎn)程存儲(chǔ)的高吞吐量訪問(wèn)；
• 提供高級(jí)、完全可組合的索引操作和虛擬視圖。

TensorStore 已被用于解決科學(xué)計(jì)算中的工程挑戰(zhàn)，還被用于創(chuàng)建大型機(jī)器學(xué)習(xí)模型，例如用來(lái)管理 PaLM 在分布式訓(xùn)練期間的模型參數(shù)（檢查點(diǎn)）。

GitHub 地址：https://github.com/google/tensorstore

用于數(shù)據(jù)訪問(wèn)和操作的 API

TensorStore 提供了一個(gè)簡(jiǎn)單的 Python API 用于加載和操作大型數(shù)組數(shù)據(jù)。例如，下面的代碼創(chuàng)建了一個(gè) TensorStore 對(duì)象，該對(duì)象代表一個(gè) 56 萬(wàn)億體素的蒼蠅大腦 3D 圖像，并允許訪問(wèn) NumPy 數(shù)組中 100x100 的圖像 patch 數(shù)據(jù)：

值得注意的是，該程序在訪問(wèn)特定的 100x100 patch 之前，不會(huì)訪問(wèn)內(nèi)存中的實(shí)際數(shù)據(jù)，因此可以加載和操作任意大的基礎(chǔ)數(shù)據(jù)集，而無(wú)需將整個(gè)數(shù)據(jù)集存儲(chǔ)在內(nèi)存中。TensorStore 使用與標(biāo)準(zhǔn) NumPy 基本相同的索引和操作語(yǔ)法。 

此外，TensorStore 還為高級(jí)索引功能提供廣泛支持，包括對(duì)齊、虛擬視圖等。

下面的代碼演示了如何使用 TensorStore 創(chuàng)建一個(gè) zarr 數(shù)組，以及 TensorStore 的異步 API 如何實(shí)現(xiàn)更高的吞吐量：

安全和性能擴(kuò)展

眾所周知，分析和處理大型數(shù)據(jù)集需要大量的計(jì)算資源，通常需要分布在多個(gè)機(jī)器上的 CPU 或加速器內(nèi)核的并行化來(lái)實(shí)現(xiàn)。因此，TensorStore 的一個(gè)基本目標(biāo)是實(shí)現(xiàn)并行處理，達(dá)到既安全又高性能的目的。事實(shí)上，在谷歌數(shù)據(jù)中心內(nèi)的測(cè)試中，他們發(fā)現(xiàn)隨著 CPU 數(shù)量的增加，TensorStore 讀寫性能幾乎呈線性增長(zhǎng)：

在谷歌云存儲(chǔ) (GCS) 上對(duì) zarr 格式數(shù)據(jù)集的讀寫性能，讀和寫性能與計(jì)算任務(wù)的數(shù)量幾乎成線性增長(zhǎng)。

TensorStore 還提供了可配置的內(nèi)存緩存和異步 API，以允許讀寫操作在程序完成其他工作時(shí)在后臺(tái)繼續(xù)執(zhí)行。為了使 TensorStore 的分布式計(jì)算與數(shù)據(jù)處理工作流兼容，谷歌還將 TensorStore 與 Apache Beam 等并行計(jì)算庫(kù)集成。

示例展示

示例 1 語(yǔ)言模型：最近一段時(shí)間，機(jī)器學(xué)習(xí)領(lǐng)域出現(xiàn)了一些 PaLM 等高級(jí)語(yǔ)言模型。這些模型包含數(shù)千億個(gè)參數(shù)，在自然語(yǔ)言理解和生成方面表現(xiàn)出驚人的能力。不過(guò)這些模型對(duì)計(jì)算設(shè)施提出了挑戰(zhàn)，特別是，訓(xùn)練一個(gè)像 PaLM 這樣的語(yǔ)言模型需要數(shù)千個(gè) TPU 并行工作。

其中有效地讀取和寫入模型參數(shù)是訓(xùn)練過(guò)程面臨的一個(gè)問(wèn)題：例如訓(xùn)練分布在不同的機(jī)器上，但參數(shù)又必須定時(shí)的保存到 checkpoint 中；又比如單個(gè)訓(xùn)練必須僅讀取特定參數(shù)集，以避免加載整個(gè)模型參數(shù)集（可能是數(shù)百 GB）所需的開(kāi)銷。

TensorStore 可以解決上述問(wèn)題。它已被用于管理大型（multipod）模型相關(guān)的 checkpoint，并已與 T5X 和 Pathways 等框架集成。TensorStore 將 Checkpoint 轉(zhuǎn)換為 zarr 格式存儲(chǔ)，并選擇塊結(jié)構(gòu)以允許每個(gè) TPU 的分區(qū)并行獨(dú)立地讀取和寫入。

當(dāng)保存 checkpoint 時(shí)，參數(shù)以 zarr 格式寫入，塊網(wǎng)格進(jìn)一步被劃分，以用于在 TPU 上劃分參數(shù)網(wǎng)格。主機(jī)為分配給該主機(jī)的 TPU 的每個(gè)分區(qū)并行寫入 zarr 塊。使用 TensorStore 的異步 API，即使數(shù)據(jù)仍被寫入持久存儲(chǔ)，訓(xùn)練也會(huì)繼續(xù)進(jìn)行。當(dāng)從 checkpoint 恢復(fù)時(shí)，每個(gè)主機(jī)只讀取分配給該主機(jī)的分區(qū)塊。

示例 2 大腦 3D 映射：突觸分辨連接組學(xué)的目標(biāo)是在單個(gè)突觸連接水平上繪制動(dòng)物和人腦的連線。完成這一目標(biāo)需要在毫米或更大的視野范圍內(nèi)以極高的分辨率 (納米級(jí)) 對(duì)大腦進(jìn)行成像，由此產(chǎn)生的數(shù)據(jù)大小達(dá)到 PB 級(jí)。然而，即使是現(xiàn)在，數(shù)據(jù)集也面臨著存儲(chǔ)、處理等方面的問(wèn)題，即使是單個(gè)大腦樣本也可能需要數(shù)百萬(wàn) GB 的空間。

谷歌已經(jīng)使用 TensorStore 來(lái)解決與大規(guī)模連接組學(xué)數(shù)據(jù)集相關(guān)的計(jì)算挑戰(zhàn)。具體而言，TensorStore 已經(jīng)開(kāi)始管理一些連接組學(xué)數(shù)據(jù)集，并將谷歌云存儲(chǔ)作為底層對(duì)象存儲(chǔ)系統(tǒng)。

目前，TensorStore 已被用于人類大腦皮層數(shù)據(jù)集 H01，原始成像數(shù)據(jù)為 1.4 PB（約為 500000 * 350000 * 5000 像素）。之后原始數(shù)據(jù)被細(xì)分為 128x128x16 像素的獨(dú)立塊，以「Neuroglancer precomputed」格式存儲(chǔ)，TensorStore 可以很容易的對(duì)其進(jìn)行操作。

利用 TensorStore 可以輕松訪問(wèn)和操作底層數(shù)據(jù)（蒼蠅大腦重建）

想要上手一試的小伙伴，可以使用以下方法安裝 TensorStore PyPI 包：

pip install tensorstore