跳到内容

使用 BioNeMo 进行开发

在本页中，我们将介绍代码库的组织结构，以及 BioNeMo 框架用户的两个主要开发工作流程（模型训练和微调）所需的设置。对于这两个工作流程，我们建议按照初始化指南中的说明设置 NGC_CLI_API_KEY 环境变量。此环境变量是脚本所必需的，该脚本将用于从 NGC 下载模型检查点和数据，以用于这些工作流程。

BioNeMo 代码概述

BioNeMo 代码库被组织为一个元包，其中包含许多 Python 包。我们这样设计 BioNeMo 是期望用户在其自己的项目中导入和使用 BioNeMo。通过以这种方式组织代码，我们确保 BioNeMo 开发者遵循与我们对最终用户期望的相似模式。

每个模型都存储在 sub-packages 的子目录中。模型示例包括

  • bionemo-esm2: ESM-2 模型
  • bionemo-geneformer: Geneformer 模型
  • bionemo-example_model: 一个最小的 MNIST 模型示例，演示了如何编写一个轻量级的 Megatron 模型，该模型实际上不支持任何 Megatron 并行性，但只要您仅使用数据并行来训练，就应该可以正常运行。

我们还包括有用的实用工具包，例如

  • bionemo-scdl: 单细胞数据加载器 (SCDL) 提供了一个数据集实现，可供 bionemo 包中的下游单细胞模型使用。
  • bionemo-testing: 一套在测试中有用的实用程序，例如 torch.testing 或 np.testing。

最后，一些包代表通用功能和抽象基类，它们公开了对与 NeMo2 交互有用的 API。这些示例包括

  • bionemo-core: 高级 API
  • bionemo-llm: 多个大型语言模型（例如 BERT 变体）共享的代码的抽象基类。

包结构

在每个 Bionemo 包中，都采用一致的结构来促进组织和可维护性。以下组件存在于每个包中

  • pyproject.toml：定义包元数据，包括版本、包名称和要安装的可执行脚本。
  • src：包含包的所有源代码。每个包都包含一个顶级的 bionemo 文件夹，该文件夹用作导入的主要命名空间。在构建过程中，所有 bionemo/* 源文件都将合并到一个包中，并将唯一的子目录名称附加到 bionemo 目录。
  • tests：存放所有包测试。测试文件的约定是将其放置在与相应的 src 文件相同的路径中，但在 tests 目录中，并在测试文件名中添加 test_ 前缀。例如，要测试模块级文件 src/bionemo/my_module，应创建一个测试文件 tests/bionemo/test_my_module.py。类似地，要测试特定文件 src/bionemo/my_module/my_file.py，测试文件应命名为 tests/bionemo/my_module/test_my_file.py。运行 py.test sub-packages/my_package 将执行 tests 目录中的所有测试。
  • examples：某些包包含一个 examples 目录，其中包含 Jupyter Notebook (.ipynb) 文件，这些文件被聚合到主文档中。
  • README.md：核心包 README 文件充当每个子包上传到 PyPI 时的主要文档。
  • LICENSE：为了保持一致性，所有 Bionemo 包都应使用 Apache-2.0 许可证。通过向 BioNeMo 贡献代码，您承认允许根据 Apache v2 许可证重新发布代码。

模型训练过程

注意

另请参阅 模型训练

从 BioNeMo 预训练模型的过程涉及运行位于 scripts 目录中的脚本。每个脚本都公开了一个命令行界面 (CLI)，其中包含并记录了该模型的可用选项。

要预训练模型，您需要使用必需的参数运行相应的脚本。例如，要预训练 ESM-2 和 Geneformer 模型，您将分别调用 train_esm2 和 train_geneformer 可执行文件。

这些脚本提供了各种可以为预训练定制的选项，例如

  • 数据目录和路径
  • 模型检查点路径
  • 用于跟踪的实验名称
  • GPU 和节点数量
  • 验证检查间隔
  • 数据集工作进程数
  • 步数
  • 序列长度
  • 微批大小
  • 验证批次限制

您可以在运行脚本时使用命令行参数指定这些选项。对于每个可用的脚本，您可以使用 --help 选项来解释该模型的可用选项。

有关预训练模型的更多信息，请参阅 ESM-2 预训练教程。

微调

模型微调过程涉及使用 download_bionemo_data 脚本下载所需的模型检查点。此脚本接受模型名称和版本作为参数，以及数据源，数据源可以是 ngc（默认）或 NVIDIA 员工的 pbss。

要查看可用资源（包括模型检查点和数据集）的列表，您可以使用以下命令

download_bionemo_data --list-resources

步骤 1：下载数据和检查点

要下载数据和检查点，请使用以下命令

export DATA_SOURCE="ngc"
MODEL_CKPT=$(download_bionemo_data <model_name>/<checkpoint_name>:<version> --source $DATA_SOURCE);

将 <model_name> 替换为所需的模型（例如，esm2 或 geneformer），将 <version> 替换为所需的版本，并将 <checkpoint_name> 替换为所需的检查点名称。

此外，您可以使用以下命令从 NGC 下载可用的数据集，进行与上述模型检查点下载命令类似的替换

TEST_DATA_DIR=$(download_bionemo_data <model_name>/testdata:<version> --source $DATA_SOURCE);

或者，您可以使用自己的数据，方法是按照初始化指南中的讨论，使用卷挂载配置您的容器运行。

步骤 2：调整训练过程

微调可能涉及指定与用于训练模型初始版本不同的模型和损失组合。微调步骤将是特定于应用程序的，但一般步骤包括

  1. 准备数据集：收集并准备您的数据集，包括序列数据和目标值。此步骤至关重要，以确保您的数据集的格式可用于训练。
  2. 创建自定义数据集类：定义一个可以处理您的特定数据格式的自定义数据集类。此类应能够初始化数据集并检索单个数据点。
  3. 创建数据模块：定义一个数据模块，用于准备训练数据。这包括数据加载、分词和批处理等任务。
  4. 微调模型：使用预训练模型作为起点，并在您的数据集上对其进行微调。这涉及调整模型的参数以适应您的特定任务和数据集。
  5. 配置微调：为微调过程设置各种超参数，例如批大小、训练步数和学习率。这些超参数会显着影响微调模型的性能。
  6. 运行推理：模型微调完成后，使用它对新的、未见过的数据进行预测。

有关微调模型的更多信息，请参阅 ESM-2 微调教程。

高级开发者文档

有关高级开发信息（例如，开发 BioNeMo 的源代码），请参阅 BioNeMo GitHub 存储库主页上的 README。