儲存和操作n維資料的難題，谷歌用一個開源軟體庫解決了

電腦科學和機器學習 (ML) 的許多應用都需要處理跨座標系的多維資料集，並且單個數據集可能也需要儲存 TB 或 PB 的資料。另一方面，使用此類資料集也具有挑戰性，因為使用者可能會以不規則的間隔和不同的規模讀取和寫入資料，通常還會執行大量的並行工作。

為了解決上述問題，谷歌開發了一個開源的 C++ 和 Python 軟體庫 TensorStore，專為儲存和操作 n 維資料而設計。谷歌 AI 負責人 Jeff Dean 也在推特上發文表示 TensorStore 現已正式開源。

TensorStore 的主要功能包括：

• 提供統一的 API 用於讀寫多種陣列格式，包括 zarr 和 N5；

• 原生支援多種儲存系統，包括谷歌雲端儲存、本地和網路檔案系統、HTTP 伺服器和記憶體儲存；

• 支援讀 / 寫快取和事務，具有很強的原子性、隔離性、一致性和永續性（ACID）特性；

• 支援從多個程序和機器進行安全、高效的併發訪問；

• 提供非同步 API 以實現對高延遲遠端儲存的高吞吐量訪問；

• 提供高階、完全可組合的索引操作和虛擬檢視。

TensorStore 已被用於解決科學計算中的工程挑戰，還被用於建立大型機器學習模型，例如用來管理 PaLM 在分散式訓練期間的模型引數（檢查點）。

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

用於資料訪問和操作的 API

TensorStore 提供了一個簡單的 Python API 用於載入和操作大型陣列資料。例如，下面的程式碼建立了一個 TensorStore 物件，該物件代表一個 56 萬億體素的蒼蠅大腦 3D 影象，並允許訪問 NumPy 陣列中 100x100 的影象 patch 資料：

值得注意的是，該程式在訪問特定的 100x100 patch 之前，不會訪問記憶體中的實際資料，因此可以載入和操作任意大的基礎資料集，而無需將整個資料集儲存在記憶體中。TensorStore 使用與標準 NumPy 基本相同的索引和操作語法。

此外，TensorStore 還為高階索引功能提供廣泛支援，包括對齊、虛擬檢視等。

下面的程式碼演示瞭如何使用 TensorStore 建立一個 zarr 陣列，以及 TensorStore 的非同步 API 如何實現更高的吞吐量：

安全和效能擴充套件

眾所周知，分析和處理大型資料集需要大量的計算資源，通常需要分佈在多個機器上的 CPU 或加速器核心的並行化來實現。因此，TensorStore 的一個基本目標是實現並行處理，達到既安全又高效能的目的。事實上，在谷歌資料中心內的測試中，他們發現隨著 CPU 數量的增加，TensorStore 讀寫效能幾乎呈線性增長：

在谷歌雲端儲存 (GCS) 上對 zarr 格式資料集的讀寫效能，讀和寫效能與計算任務的數量幾乎成線性增長。

TensorStore 還提供了可配置的記憶體快取和非同步 API，以允許讀寫操作在程式完成其他工作時在後臺繼續執行。為了使 TensorStore 的分散式計算與資料處理工作流相容，谷歌還將 TensorStore 與 Apache Beam 等平行計算庫整合。

示例展示

示例 1 語言模型：最近一段時間，機器學習領域出現了一些 PaLM 等高階語言模型。這些模型包含數千億個引數，在自然語言理解和生成方面表現出驚人的能力。不過這些模型對計算設施提出了挑戰，特別是，訓練一個像 PaLM 這樣的語言模型需要數千個 TPU 並行工作。

其中有效地讀取和寫入模型引數是訓練過程面臨的一個問題：例如訓練分佈在不同的機器上，但引數又必須定時的儲存到 checkpoint 中；又比如單個訓練必須僅讀取特定引數集，以避免載入整個模型引數集（可能是數百 GB）所需的開銷。

TensorStore 可以解決上述問題。它已被用於管理大型（multipod）模型相關的 checkpoint，並已與 T5X 和 Pathways 等框架整合。TensorStore 將 Checkpoint 轉換為 zarr 格式儲存，並選擇塊結構以允許每個 TPU 的分割槽並行獨立地讀取和寫入。

當儲存 checkpoint 時，引數以 zarr 格式寫入，塊網格進一步被劃分，以用於在 TPU 上劃分引數網格。主機為分配給該主機的 TPU 的每個分割槽並行寫入 zarr 塊。使用 TensorStore 的非同步 API，即使資料仍被寫入持久儲存，訓練也會繼續進行。當從 checkpoint 恢復時，每個主機只讀取分配給該主機的分割槽塊。

示例 2 大腦 3D 對映：突觸分辨連線組學的目標是在單個突觸連線水平上繪製動物和人腦的連線。完成這一目標需要在毫米或更大的視野範圍內以極高的解析度 (奈米級) 對大腦進行成像，由此產生的資料大小達到 PB 級。然而，即使是現在，資料集也面臨著儲存、處理等方面的問題，即使是單個大腦樣本也可能需要數百萬 GB 的空間。

谷歌已經使用 TensorStore 來解決與大規模連線組學資料集相關的計算挑戰。具體而言，TensorStore 已經開始管理一些連線組學資料集，並將谷歌雲端儲存作為底層物件儲存系統。

目前，TensorStore 已被用於人類大腦皮層資料集 H01，原始成像資料為 1.4 PB（約為 500000 * 350000 * 5000 畫素）。之後原始資料被細分為 128x128x16 畫素的獨立塊，以「Neuroglancer precomputed」格式儲存，TensorStore 可以很容易的對其進行操作。

利用 TensorStore 可以輕鬆訪問和操作底層資料（蒼蠅大腦重建）

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

pip install tensorstore

原文連結：

http://ai.googleblog.com/2022/09/tensorstore-for-high-performance.html