在 HSSM 中使用存储在 Hugging Face 的自定义 ONNX 模型

本指南将向您介绍如何在 HSSM（分层状态空间模型）框架中使用存储在 Hugging Face 的自定义 ONNX 模型。

前提条件

1. Python 3.8 或更高版本。
2. 在您的 Python 环境中安装了 HSSM 库。
3. 在 Hugging Face 模型中心存储了一个预训练的 ONNX 模型。

分步指南

步骤 1：导入必要的库

import pandas as pd
import hssm
import ssms.basic_simulators

pytensor.config.floatX = \"float32\"

步骤 2：定义 HSSM 配置

您需要定义模型的配置。确保将 log-likelihood 类型定义为 "approx_differentiable"，并在 loglik 字段中提供 Hugging Face 模型名称。

my_hssm = hssm.HSSM(
        data=dataset_lan,
        loglik_kind = \"approx_differentiable\",
        loglik = \"levy.onnx\",
        model=\"custom\",
        model_config= {
            \"backend\": \"jax\",
            \"list_params\": [\"v\", \"a\", \"z\", \"alpha\", \"t\"],
            \"bounds\": {
                \"v\": (-3.0, 3.0),
                \"a\": (0.3, 3.0),
                \"z\": (0.1, 0.9),
                \"alpha\": (1.0, 2.0),
                \"t\": (1e-3, 2.0),
            },
        }                  
    )

这将使用来自 Hugging Face 仓库的自定义 ONNX 模型 levy.onnx 创建一个 HSSM 对象 my_hssm。

my_hssm.sample(cores=2, draws=500, tune=500, mp_ctx=\"forkserver\")

将 ONNX 文件上传到 Hugging Face 仓库

如果您的 ONNX 文件目前不在您的 Hugging Face 仓库中，您可以按照以下步骤将其添加进去：

1. 从 huggingface_hub 导入 HfApi 模块：

from huggingface_hub import HfApi

2. 使用 upload_file 方法上传 ONNX 文件：

api = HfApi()
api.upload_file(
    path_or_fileobj=\"test.onnx\",
    path_in_repo=\"test.onnx\",
    repo_id=\"franklab/HSSM\",
    repo_type=\"model\",
    create_pr=True,
)

执行这些步骤将在 Hugging Face 上生成一个拉取请求（PR），该请求随后将由我们团队的成员进行审核。

创建拉取请求和新的 ONNX 模型

1. 在 Hugging Face 上创建拉取请求

   访问以下链接：Hugging Face PR

   这样做将在 Hugging Face 上生成一个拉取请求，该请求将由我们的团队成员进行审核。

2. 创建自定义 ONNX 模型

   在 PyTorch 中建立网络配置和状态字典文件

   要构建自定义模型并将其保存为 ONNX 文件，您必须在 PyTorch 中创建网络配置文件和状态字典文件。请参阅 LANFactory 包 的 README 中概述的说明。

   将网络配置和状态字典文件转换为 ONNX

   生成网络配置和状态字典文件后，您需要将这些文件转换为 ONNX 格式。