跳到内容

BioNemo-SCDL：用于单细胞基础模型可扩展训练的单细胞数据加载。

软件包概述

BioNeMo-SCDL 为单细胞数据提供了一个独立的 pytorch 兼容数据集类，并具有一致的 API。BioNeMo-SCDL 由 NVIDIA 开发和维护。此软件包可以独立于 BioNeMo 运行。它在以下方面改进了基于简单 AnnData 的数据集类：

  • 跨输入格式的一致 API，并承诺在软件包版本之间保持一致。
  • 加载大型数据集时性能得到提升。它允许加载和快速迭代大型数据集。
  • 能够使用远超内存的数据集。这是因为数据集以 numpy 内存映射格式存储。
  • 此外，还将大型（远超内存）AnnData 文件转换为 SCDL 格式。
  • [未来] 完全支持参差不齐的数组（即，具有不同特征计数的数据集；目前只有部分 API 功能支持参差不齐的数组）。
  • [未来] 支持改进的压缩。

BioNeMo-SCDL 的 API 类似于 AnnData，因此代码更改非常少。在大多数情况下，只需将属性简单地替换为函数即可（即，将 data.n_obs 替换为 data.number_of_rows()）。

安装

此软件包可以通过以下方式安装

pip install bionemo-scdl

使用方法

获取示例数据

以下是如何处理来自 CellxGene 的约 25,000 个细胞的示例数据集

下载 "https://datasets.cellxgene.cziscience.com/97e96fb1-8caf-4f08-9174-27308eabd4ea.h5ad" 到 hdf5s/97e96fb1-8caf-4f08-9174-27308eabd4ea.h5ad

从 H5AD 文件加载单细胞数据集

from bionemo.scdl.io.single_cell_memmap_dataset import SingleCellMemMapDataset

data = SingleCellMemMapDataset("97e_scmm", "hdf5s/97e96fb1-8caf-4f08-9174-27308eabd4ea.h5ad")

这将创建一个 SingleCellMemMapDataset，它存储在 97e_scmm 中，以大型内存映射数组形式存储，从而可以快速访问大于系统可用 RAM 容量的数据集。

如果数据集很大，则可以延迟加载 AnnData 文件，然后根据分页方式的行块读取。这可以通过在实例化 SingleCellMemMapDataset 时设置参数来完成： - paginated_load_cutoff，它设置 AnnData 文件以分页方式读取的最小文件大小（以 MB 为单位）。 - load_block_row_size，它是每次读入内存的行数。

查询单细胞数据集并探索 API

data.number_of_rows()
## 25382

data.number_of_variables()
## [34455]

data.number_of_values()
## 874536810

data.number_nonzero_values()
## 26947275

将 SCDL (单细胞数据加载器) 数据集保存到磁盘

当您打开 SCDL 数据集时，您必须选择一个路径来存储后备数据结构。但是，不能保证这些结构在运行时处于有效的序列化状态。

调用 save 方法可保证磁盘上的对象处于有效的序列化状态，此时当前 python 进程可以退出，并且该对象稍后可以由另一个进程加载。

data.save()

从 SCDL 存档加载 SCDL 数据集

当您准备好重新加载 SCDL 数据集时，只需传递序列化数据的路径即可

reloaded_data = SingleCellMemMapDataset("97e_scmm")

在模型训练中使用 SCDL 数据集

SCDL 实现了 PyTorch Dataset 抽象类的必需函数。您可以使用 PyTorch 兼容的 DataLoaders 从 SCDL 类加载数据批次。当批次大小为 1 时，这可以在没有整理函数的情况下运行。当批次大小大于 1 时，有一个整理函数 (collate_sparse_matrix_batch)，它会将多个稀疏数组整理为 CSR（压缩稀疏行）torch 张量格式。

from torch.utils.data import DataLoader
from bionemo.scdl.util.torch_dataloader_utils import collate_sparse_matrix_batch

## Mock model: you can remove this and pass the batch to your own model in actual code.
model = lambda x : x

dataloader = DataLoader(data, batch_size=8, shuffle=True, collate_fn=collate_sparse_matrix_batch)
n_epochs = 2
for e in range(n_epochs):
    for batch in dataloader:
        model(batch)

示例

examples 目录包含利用 SCDL 的各种示例。

将现有的 Cell x Gene 数据转换为 SCDL

如果有多个 AnnData 文件，它们可以转换为单个 SingleCellMemMapDataset。如果 hdf5 目录有一个或多个 AnnData 文件，则 SingleCellCollection 类会爬取文件系统以递归查找 AnnData 文件（扩展名为 h5ad）。

要转换现有的 AnnData 文件，您可以编写自己的脚本使用 SCDL API，或者使用便捷脚本 convert_h5ad_to_scdl。

这是一个例子

convert_h5ad_to_scdl --data-path hdf5s --save-path example_dataset

未来工作和路线图

SCDL 目前处于公开 Beta 测试阶段。未来，预计数据压缩和数据加载性能将得到改进。

许可证

BioNemo-SCDL 具有 Apache 2.0 许可证，如 LICENSE 文件中所示。