Существует множество интересных современных приложений информатики и машинного обучения, которые включают в себя манипулирование многомерными наборами данных, охватывающими одну большую систему координат. Например, моделирование погоды на основе атмосферных измерений по пространственной сетке или прогнозирование медицинских изображений на основе значений интенсивности многоканального изображения в 2D- или 3D-сканировании. При таких настройках даже для одного набора данных может потребоваться терабайт или петабайт хранилища данных. С такими наборами данных также сложно работать, поскольку пользователи могут считывать и записывать данные с нерегулярными интервалами и в различных масштабах и часто заинтересованы в выполнении анализа с использованием множества машин, работающих параллельно. Google AI представляет TensorStore, программную библиотеку C++ и Python с открытым исходным кодом, предназначенную для хранения и обработки n-мерных данных.

TensorStore

Это библиотека с открытым исходным кодом, которая предоставляет единый API для чтения и записи нескольких форматов массивов, включая zarr и N5. Он изначально поддерживает несколько систем хранения, включая Google Cloud Storage, локальные и сетевые файловые системы, HTTP-серверы и хранилище в оперативной памяти. Он также поддерживает кэширование чтения/записи и транзакции с надежными гарантиями ACID. Кроме того, он предлагает асинхронный API для обеспечения высокоскоростного доступа даже к удаленному хранилищу с высокой задержкой. Наконец, он предоставляет расширенные, полностью компонуемые операции индексирования и виртуальные представления.

TensorStore использовался для решения различных инженерных задач, таких как управление и обработка больших наборов данных в нейробиологии, а также создание крупномасштабных моделей машинного обучения. Примером последнего является PaLM, который решает проблему управления параметрами модели во время распределенного обучения.

API для доступа к данным и управления ими

TensorStore Python API упрощает загрузку больших массивов данных и управление ими. В приведенном ниже примере создается объект TensorStore, представляющий 56 триллионов вокселей 3D изображения мозга мухи, и доступ к небольшому фрагменту 100x100 данных в виде массива NumPy.

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])

В следующем примере показано, как можно использовать 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()

Безопасное и эффективное масштабирование

TensorStore предназначен для обеспечения параллельной обработки отдельных наборов данных, которая является одновременно безопасной и высокопроизводительной. В ходе тестирования в центрах обработки данных Google было обнаружено, что TensorStore имеет почти линейное масштабирование производительности чтения и записи по мере увеличения количества процессоров.

TensorStore обеспечивает высокую производительность за счет реализации основных операций на C++, многопоточности для таких операций, как кодирование/декодирование и сетевой ввод-вывод, а также разделения больших наборов данных на более мелкие единицы. Кэширование в памяти и асинхронный API помогают еще больше повысить производительность, сокращая количество взаимодействий с более медленными системами хранения и позволяя выполнять фоновые операции.

Варианты использования TensorStore

Ниже приведены два варианта использования TensorStore:

Языковые модели

Более продвинутые языковые модели, такие как PaLM, являются интересными недавними разработками в области машинного обучения. Эти нейронные сети содержат сотни миллиардов параметров и демонстрируют удивительные возможности в понимании и генерации естественного языка. Эти модели также расширяют границы вычислительной инфраструктуры; в частности, для обучения такой языковой модели, как PaLM, требуется, чтобы тысячи TPU работали параллельно.

Проблема, возникающая в процессе обучения, заключается в быстром и эффективном чтении и записи параметров модели. Обучение проводится на множестве разных машин, но параметры нужно регулярно сохранять в одном месте («контрольной точке») в системе хранения, которую нельзя удалить. Кроме того, каждое отдельное задание обучения должно иметь возможность считывать только те параметры, которые ему необходимы, чтобы избежать дополнительного времени, необходимого для загрузки всего набора параметров (который может составлять сотни гигабайт).

TensorStore — это инструмент, который использовался для управления контрольными точками, связанными с крупномасштабными моделями, обученными с помощью JAX. Он был интегрирован с такими платформами, как T5X и Pathways. Модельный параллелизм используется для разделения полного набора параметров, который может занимать более терабайта памяти, на сотни TPU. Контрольные точки хранятся в формате zarr с использованием TensorStore со структурой фрагментов, выбранной таким образом, чтобы раздел для каждого TPU мог читаться и записываться независимо параллельно.

3D-картирование мозга

Область коннектомики разрешения синапсов посвящена картированию проводки мозга на невероятно детальном уровне, вплоть до отдельных синаптических соединений. Для этого требуется визуализация мозга с очень высоким разрешением (нанометры) на больших площадях (миллиметры и более). Это может создавать наборы данных размером в петабайты. В будущем, когда ученые планируют составить карту всего мозга мыши или примата, эти наборы данных потенциально могут достигать эксабайтов. Однако даже текущие наборы данных чрезвычайно велики и создают серьезные проблемы с точки зрения хранения, обработки и обработки. Например, для одного образца мозга могут потребоваться миллионы гигабайт дискового пространства с системой координат (пиксельным пространством) из сотен тысяч пикселей в каждом измерении.

TensorStore используется для решения вычислительных задач, связанных с крупномасштабными наборами коннектомных данных. В частности, TensorStore управляет некоторыми из крупнейших и наиболее широко используемых коннектомных наборов данных, используя Google Cloud Storage в качестве базовой системы хранения объектов.

Заключение

Чтобы начать использовать TensorStore Python API, вы можете установить пакет tensorstore PyPI, используя:

pip install tensorstore

Подробности использования см. в руководствах и документации по API. Другие варианты установки и использование C++ API см. в инструкциях по установке.

Чтобы узнать больше интересных и захватывающих блогов, следите за обновлениями!

Следуй за мной: М. Хасиб Хассан