使用Python进行基于仿真的推理(SBI)

SBI是一个用于进行基于仿真的推理(Simulation-Based Inference,SBI)的Python库。该库基于对给定模型在不同参数组下的多次仿真,推理输入参数分布与输出之间的概率关系,通过预训练神经网络模型,在新的输入抵达时快速给出参数估计结果。

SBI的安装

SBI可以使用pip命令进行安装:

pip install sbi

SBI依赖PyTorch,使用pip命令安装SBI时默认会安装CPU版本的PyTorch,等价于在安装SBI前执行:

pip install torch torchvision

针对适用于不同操作系统平台,以及CUDA等并行计算框架的PyTorch,请在安装SBI前安装,请参考:https://pytorch.org/get-started/locally/

例如,若需要PyTorch可在Windows环境上调用CUDA 12.6,需在安装SBI前执行:

pip install torch torchvision --index-url *https://download.pytorch.org/whl/cu126

基本流程

根据贝叶斯公式:

# Bayes' Theorem
#     Param - Model parameters
#     Observed - Observed data
p(Param | Observed) = p(Observed | Param) * p(Param) / p(Observed)

故SBI的流程为:

  1. 确定先验分布(确定p(Param)
  2. 采样多个仿真数据,进行多轮仿真,得到多个(Param, Observed)的输入-输出对
  3. 训练用于拟合后验概率p(Param | Observed)、似然p(Observed | Param)或其它变体的神经网络(神经密度估计器或神经分类器),常用算法包括:
    • NPE(Neural Posterior Estimation):训练一个深度神经密度估计器来直接近似后验分布p(Param | Observed)
    • NLE(Neural Likelihood Estimation):训练一个神经网络,使用条件密度估计器(归一化流)来近似似然p(Observed | Param)
    • NRE(Neural Ratio Estimation):训练一个神经分类器,通过区分来自联合分布p(Param, Observed)的样本和来自边缘分布p(Param)p(Observed)的样本,来估计似然与证据的比率r(Observed | Param) = p(Observed | Param) / p(Observed)
  4. 从训练完毕的神经网络建立后验概率p(Param | Observed)采样器

定义仿真器

仿真器(Simulator),即模型,是对真实过程的虚拟表达。其接收一个一维参数数组arrParams以及潜在的扩展参数,并传回包含了仿真结果或其特征(低维表达)的数组arrFeatures

例如用下面的仿真器SBIParamToFeatures()表示一个过程,第一个参数arrParams接受一个包含两个元素的一维数组,第二个参数IsResultUseNumPy控制返回类型为numpy.array还是torch.Tensor。其返回值为一个包含两个元素的一维数组,系对两个参数的简单修改,模拟一种数据特征:

# NumPy
import numpy as np
 
# PyTorch
import torch
 
# Simulation-Based Inference (SBI)
import sbi
import sbi.utils
import sbi.inference
 
# Object process simulator
# This function simulates a real-world process, or a wrapper of another simulator
# This function accepts a parameter array arrParams with 2 elements [dParam1, dParam2]
# And returns a low-level representation (features) of [dFeature1, dFeature2], where:
#     dFeature1 = dParam2 + 1
#     dFeature2 = dParam1 + 2
# Parameter IsResultUseNumPy selects the return value type of this function (tensor or numpy array)
def SBIParamToFeatures(arrParams : list, IsResultUseNumPy : bool) -> list :
    arrFeatures = np.array([arrParams[1] + 1, arrParams[0] + 2])
    if IsResultUseNumPy :
        return arrFeatures
    else :
        return torch.as_tensor(arrFeatures)
    #End If
#End Function

定义参数分布与边界的先验

首先需要定义模型参数的先验分布p(Param)。可以使用sbi.utils.BoxUniform对象定义参数的边界,该类将参数预设为均匀分布(Uniform Distribution)。使用sbi.utils.process_prior()函数确保先验数据格式满足要求:

# Main entry point
if __name__ == "__main__" :
    # Ref. *https://sbi.readthedocs.io/en/latest/tutorials/Example_00_HodgkinHuxleyModel.html
    # Bayes' Theorem
    #     p(Param | Observed) = p(Observed | Param) * p(Param) / p(Observed)
    #         Param - Model parameters
    #         Observed - Observed data
 
    # Defining priori parameter bounds and distribution (build prior distribution p(Param))
    #     dParam1 in [1,10]
    #     dParam2 in [25,45]
    arrOptimizationBounds = [
        (1,10),
        (25,45)
    ]
    arrSBILowerBounds = torch.as_tensor([Bound[0] for Bound in arrOptimizationBounds])
    arrSBIUpperBounds = torch.as_tensor([Bound[1] for Bound in arrOptimizationBounds])
    bndSBIUniformedBounds = sbi.utils.BoxUniform(low=arrSBILowerBounds, high=arrSBIUpperBounds)
    dstSBIPriorDistribution, nSBIParam, IsSBIPriorNumPy = sbi.utils.process_prior(bndSBIUniformedBounds)

预处理仿真器函数

当前定义的仿真器函数SBIParamToFeatures()一次只能输入一组参数,并且包含一个扩展参数IsResultUseNumPy,而SBI对仿真器有如下要求:
1. 函数只接受一个参数,即函数签名形如def SBIParamToFeatures(arrParams)
2. 支持批量参数输入与仿真。即仿真器参数arrParams的结构为(nSimCount, nParamCount),其中nSimCount为本次调用需进行的仿真次数,nParamCount为单次仿真所需的参数量。
3. 支持批量返回仿真结果。即仿真器返回值arrFeatures的结构为(nSimCount, nFeatureCount),其中nSimCount为本次调用需进行的仿真次数,nFeatureCount为单次仿真返回的结果向量长度(单次仿真特征数)。

针对第一个要求,可以利用functools.partial对象建立偏函数,将SBIParamToFeatures()除第一个参数外的参数固定:

    # Preprocessing simulator
    # Create partial version of SBIParamToFeatures()
    # Since the batched simulation wrapper of SBI (sbi.inference.simulate_for_sbi) only allows a simulator with 1 argument
    # Or you can write your own batched simulation wrapper, which accepts a parameter group arrParams returns a tuple (arrParams, arrFeatures)
    import functools
    fncSBIParamToFeaturesPartial = functools.partial(SBIParamToFeatures,
        IsResultUseNumPy=False)

针对第二和第三个要求,SBI提供了sbi.utils.process_simulator()函数,建立支持批量仿真的仿真器包裹函数(Wrapper)。该函数会自动检测给定的仿真函数是否已支持批量处理:

    # Convert to SBI function
    fncSBIParamToFeaturesWrapped = sbi.utils.process_simulator(fncSBIParamToFeaturesPartial, dstSBIPriorDistribution, IsSBIPriorNumPy)

运行仿真器

SBI提供了sbi.inference.simulate_for_sbi()函数以运行仿真器。该函数支持并行化、数据分批和进度显示等功能:

    # Defining sample data information & create sample input if necessary
    nSBISamples = 500
    print(f"Sampling {nSBISamples} samples, this may take a long time...")
 
    # Run simulator
    arrParams, arrFeatures = sbi.inference.simulate_for_sbi(simulator=fncSBIParamToFeaturesWrapped, proposal=dstSBIPriorDistribution, 
        num_simulations=nSBISamples, num_workers=1, 
        simulation_batch_size=1, seed=245, show_progress_bar=True)

扩展:自定义批量仿真

用户也可以自行为仿真器编写包裹函数,实现批量仿真和并行化等功能,以及更自由的参数定义。只要函数可以返回相应的参数序列arrParams(结构为(nSimCount, nParamCount),数据类型为Tensor或可转换为Tensor)和结果序列arrFeatures(结构为(nSimCount, nFeatureCount),数据类型为Tensor或可转换为Tensor)即可。需要注意返回的Tensor的数据类型必须为float32

arrResult = torch.as_array(arrRawResult, dtype=torch.float32)

BoxUniform提供了sample_n()sample()函数,用于产生参数序列arrParams。对于调用BoxUniform.sample_n(nSimCount)BoxUniform.sample((nSimCount,)),返回值的结构为(nSimCount, nParamCount),数据类型为Tensor。您也可以使用其它准蒙特卡罗(Quasi-Monte Carlo,QMC)采样器进行采样,例如拉丁超立方采样(Latin Hypercube Sampling,LHS),以更好地覆盖参数取值范围,但需要注意将采样结果的数据结构转换。

例如,若定义SBIParamToFeaturesBatch()函数进行批量仿真(该函数返回两个numpy.array,即arrParamsarrFeatures),则可以使用类似下面的代码:

    # Defining sample data information & create sample input if necessary
    nSBISamples = 500
    print(f"Sampling {nSBISamples} samples, this may take a long time...")
    # sample_n is deprecated
    #arrSBISimParams = dstSBIPriorDistribution.sample_n(nSBISamples)
    # You can use this
    arrSBISimParams = dstSBIPriorDistribution.sample((nSBISamples,))
    # Or you can use methods like Latin Hypercube Sampling (LHS) for better parameter space coverage
    # See: https://docs.scipy.org/doc/scipy/reference/stats.qmc.html and https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.qmc.LatinHypercube.html
    from scipy.stats import qmc
    smpParamSpaceSampler = qmc.LatinHypercube(d=nOptimizationParametersVar)
    arrParamSpaceSamplingResult = smpParamSpaceSampler.random(n=nSBISamples)
    arrSBISimParams = qmc.scale(arrParamSpaceSamplingResult, arrSBILowerBounds, arrSBIUpperBounds)
    arrSBISimParams = torch.as_tensor(np.array(arrSBISimParams), dtype=torch.float32)
 
    # Run simulator
    arrParams, arrFeatures = SBIParamToFeaturesBatch(arrSBISimParams.numpy(), 
        #Extended calling parameters here
        )
    arrParams = torch.as_tensor(arrParams)
    arrFeatures = torch.as_tensor(arrFeatures)
    # Ensure type of arrParams and arrFeatures are torch.float32 (torch.FloatTensor)
    if arrParams.type() != "torch.FloatTensor" :
        SharedUtil.LogInfo(f"Type of parameter array is {arrParams.type()}, torch.float32 (torch.FloatTensor) required")
        arrParams = arrParams.type(torch.float32)
    #End If
    if arrFeatures.type() != "torch.FloatTensor" :
        SharedUtil.LogInfo(f"Type of simulated data array is {arrFeatures.type()}, torch.float32 (torch.FloatTensor) required")
        arrFeatures = arrFeatures.type(torch.float32)
    #End If

自定义批量仿真函数在一些情况下可能是有益的,例如具有很多处理器核心的Windows平台(SBI的CPU并行计算后端可能存在并行数限制)。请注意,如果您只需要PyTorch进行仿真器输出结果的数据格式转换,建议将该转换步骤放在批量仿真函数中,而不是由每一个子仿真单独转换,因为这可能占用大量内存资源。

训练用于估计正向分布的神经网络模型并建立后验分布

为了建立由观测数据推断参数的后验分布p(Param | Observed),需要训练一个神经网络进行拟合。SBI支持多种拟合算法,请参考:https://sbi.readthedocs.io/en/latest/how_to_guide/06_choosing_inference_method.htmlhttps://sbi.readthedocs.io/en/latest/api_reference/training.html

  • SNPE:Sequential Neural Posterior Estimation
    • NPE 是唯一直接估计后验的方法。这意味着它在训练后不需要进一步的采样步骤(如 MCMC)。这使得 NPE 在训练后从后验估计中进行采样的速度最快。此外,它可以使用嵌入网络来学习汇总统计量(与 NRE 一样,但与 NLE 不同)。
    • NPE_A: Fast epsilon-free Inference of Simulation Models with Bayesian Conditional Density Estimation, Papamakarios et al., NeurIPS 2016. https://arxiv.org/abs/1605.06376
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NPE_A.html
      • 与所有 NPE 方法一样,该方法训练一个深度神经密度估计器来直接近似后验分布。同样,与所有其他 NPE 方法一样,在第一轮中,该密度估计器使用最大似然损失进行训练。
      • 该类实现了 NPE-A。NPE-A 使用最大似然损失进行多轮训练。这会导致训练收敛于提议后验而非真实后验。为了纠正这一点,SNPE-A 在训练后应用了事后校正。该校正是以解析方式进行的,并且需要使用高斯混合密度估计器。
      • 在多轮 SNPE-A 中,MoG 成分的数量随每一轮呈乘法增长:如果提议分布有 L 个成分且密度估计器有 K 个成分,那么修正后的后验分布将有 L×K 个成分。对于需要进行多轮训练的情况,建议改用 SNPE-C (APT),它能更高效地处理多轮推理。
    • NPE_B: Flexible statistical inference for mechanistic models of neural dynamics, Lueckmann, Gonçalves et al., NeurIPS 2017. https://arxiv.org/abs/1711.01861
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NPE_B.html
      • NPE-B(也称为 SNPE-B)使用重要性加权损失训练神经网络,以直接近似后验分布p(Param | Observed)。与 NPE-A 不同,这种重要性加权确保了在多轮推理中能够收敛到真实的后验分布,并且不局限于高斯提议分布。NPE-B 可以使用诸如归一化流等灵活的密度估计器。
      • 对于单轮推理,NPE-A、NPE-B 和 NPE-C 是等价的,并且均使用普通的 NLL(负对数似然)损失。
    • NPE_C: Automatic Posterior Transformation for Likelihood-free Inference, Greenberg et al., ICML 2019, https://arxiv.org/abs/1905.07488.
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NPE_C.html
      • NPE-C(也称为自动后验变换(Automatic Posterior Transformation,APT),又名 SNPE-C)通过多轮训练神经网络,以直接近似特定观测值 x_o 的后验分布。在第一轮中,NPE-C 与其他 NPE 方法等价,并且是完全摊销的(可对任何新观测值进行直接推理)。在第一轮之后,NPE-C 会根据所选的密度估计器在两种损失变体之间自动进行选择:一种是稳定且能避免泄漏的非原子损失(适用于高斯混合分布),另一种是更灵活但可能存在泄漏问题的原子损失(适用于流模型)。
    • 对于单轮推理,NPE-A、NPE-B 和 NPE-C 是等价的,并且均使用普通的 NLL(负对数似然)损失。
  • SNLE:Sequential Neural Likelihood Estimation
    • NLE 学习似然。这使得 NLE 在训练后能够模拟仿真器。为了从后验中进行采样,NLE 需要与 MCMC 或变分推断结合使用,这使得其从后验采样较慢。NLE 的一个优点是,如果有多个独立同分布(Independent Identically Distributed,IID)的观测值,那么它的仿真效率可能比 NPE 高得多。
    • NLE_A: Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows, Papamakarios et al., AISTATS 2019, https://arxiv.org/abs/1805.07226
  • SNRE:Sequential Neural Ratio Estimation
    • NRE 学习似然与证据的比率。它仅训练一个分类器(不同于训练生成模型的 NPE 和 NLE),这使得它的训练速度最快。与 NLE 一样,NRE 必须与 MCMC 或变分推断结合才能进行采样。与 NPE 一样,NRE 可以使用嵌入网络来学习汇总统计量。
    • NRE_A: Likelihood-free MCMC with Amortized Approximate Likelihood Ratios, Hermans et al., ICML 2020, https://arxiv.org/abs/1903.04057
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NRE_A.html
      • NRE-A 训练一个神经分类器,通过区分来自联合分布p(Param, Observed)的样本和来自边缘分布p(Param)p(Observed)的样本,来估计似然与证据的比率r(Observed | Param) = p(Observed | Param) / p(Observed)。随后,利用估计出的比率,通过 MCMC、拒绝采样或变分推断来进行后验采样。
      • NRE 可以进行多轮运行而无需校正,但在每一轮中都需要进行可能计算代价高昂的后验采样。
    • NRE_B: On Contrastive Learning for Likelihood-free Inference, Durkan et al., ICML 2020, https://arxiv.org/pdf/2002.03712
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NRE_B.html
      • NRE-B 是 NRE-A 的扩展,它使用对比(1-out-of-K)损失训练神经分类器,以估计似然与证据的比率。它不采用二分类,而是将来自联合分布p(Param, Observed)的一个样本与来自边缘分布p(Param)p(Observed)K-1个样本进行对比。与 NRE-A 相比,这种多类别的构建方式提高了训练的稳定性。
      • NRE 可以进行多轮运行而无需校正,但在每一轮中都需要进行可能计算代价高昂的后验采样。
    • NRC_C: Contrastive Neural Ratio Estimation, Benjamin Kurt Miller, et al. NeurIPS 2022, https://arxiv.org/abs/2210.06170
      • https://sbi.readthedocs.io/en/latest/api_reference/_autosummary/sbi.inference.NRE_C.html
      • NRE-C 使用一种“多类 Sigmoid”损失对 NRE-A 和 NRE-B 进行了推广,该损失确保在第一轮达到最优时,所估计的比率p(Param, Observed)/(p(Param)p(Observed))是精确的。这解决了 NRE-B 的比率仅定义到一个关于Observed的任意函数程度的问题。NRE-C 提供了更准确的比率估计,同时保持了对比学习的优势。
  • SBI的总体建议:
    • 如果模拟输出是高维数据(图像、时间序列等),请使用 NPE 或 NRE。
    • 如果需要快速采样,请使用 NPE。
    • 如果多个试验中存在独立同分布(IID)的观测值(且模拟成本相对较高),请使用 NRE 或 NLE。
    • 如果需要对大量观测值进行推断,请使用摊销法(这是 sbi 包的默认设置)。
    • 若需对少量观测值(<10)进行推断,请先使用摊销方法,仅当模拟成本过高时才切换至序贯方法。
    • 除非您是专家用户且清楚自己在做什么,否则请勿使用 SNPE_A 或 SNPE_B。

这里以NPE_C模型为例。训练时,使用append_simulations()函数输入训练所需的仿真参数-模型输出对,再调用train()函数训练神经网络,最后使用build_posterior()函数从神经密度估计器建立后验分布:

    # Train a neural network (neural density estimator, or neural classifier) to approximate p(Param | Observed), likelihood p(Observed | Param), or other variants
    # For estimating methods, see: https://sbi.readthedocs.io/en/latest/how_to_guide/06_choosing_inference_method.html and https://sbi.readthedocs.io/en/latest/api_reference/training.html
    #     NPE (Neural Posterior Estimation): Trains a deep neural density estimator to directly approximate the posterior p(Param | Observed).
    #         NPE is the only method which estimates the posterior directly. This means that it will not require further sampling steps (such as MCMC) after training. This makes NPE the fastest to sample from the posterior estimate after training. In addition, it can use an embedding network to learn summary statistics (just like NRE, but unlike NLE).
    #     NLE (Neural Likelihood Estimation): Trains a neural network to approximate the likelihood p(Observed | Param) using a conditional density estimator (normalizing flow).
    #         NLE learns the likelihood. This allows NLE to emulate the simulator after training. To sample from the posterior, NLE is combined with MCMC or variational inference, which makes it slower to sample. An advantage of NLE is that, if you have multiple independent identically distributed (IID) observations, then it can be much more simulation efficient than NPE
    #     NRE (Neural Ratio Estimation): Trains a neural classifier to estimate the likelihood-to-evidence ratio r(Param, Observed) = p(Observed | Param) / p(Observed) by distinguishing between samples from the joint distribution p(Param, Observed) and samples from the marginals p(Param)p(Observed). Posterior sampling is then performed via MCMC, rejection sampling, or variational inference using the estimated ratio.
    #         NRE learns the likelihood-to-evidence ratio. It trains only a classifier (unlike NPE and NLE, which train a generative model), which makes it fastest to train. Like NLE, NRE has to be combined MCMC or variational inference to sample. Like NPE, NRE can use an embedding network to learn summary statistics.
    # SBI explicit recommendations:
    #     If have high-dimensional simulation-outputs (images, time-series,…), use NPE or NRE.
    #     If you want fast sampling, use NPE.
    #     If you have IID observations across many trials (and simulations are relatively expensive), use NRE or NLE.
    #     If you want to perform inference for many observations, use amortized methods (which is the default of the sbi package).
    #     If you want to perform inference for few observations (<10), then start with amortized methods and switch to sequential methods only if the simulation cost is too high.
    #     Do not use SNPE_A or SNPE_B unless you are an expert user and know what you are doing.
    trnEstimatorTrainer = sbi.inference.NPE_C(prior=dstSBIPriorDistribution)
    trnEstimatorTrainer = trnEstimatorTrainer.append_simulations(arrParams, arrFeatures)
    estEstimator = trnEstimatorTrainer.train()
 
    # Build posterior p(Param | Observed) sampler from the trained neural network
    # For sampler selection, see: *https://sbi.readthedocs.io/en/latest/how_to_guide/09_sampler_interface.html
    #     MCMCPosterior: very accurate, but can be slow if you have many parameters.
    #     VIPosterior: can be much faster if you have many parameter. May be inaccurate. You may need to train VIPosterior before sampling.
    #     RejectionPosterior: accurate and fast, but only for very few parameters (typically less than 3).
    #     ImportanceSamplingPosterior: typically inaccurate, but can be very useful to improve the accuracy of a VIPosterior (see above).
    #     For NPE:
    #         use the DirectPosterior. If you have a likelihood available, you can refine the posterior with importance sampling
    #     For NLE or NRE:
    #         If you have very few parameters (<3), use RejectionPosterior or MCMCPosterior
    #         If you have a medium number of parameters (3-10), use MCMCPosterior
    #         If you have many parameters (>10) and the MCMCPosterior is too slow, use the VIPosterior. Optionally combine the VIPosterior with an ImportanceSamplingPosterior to improve its accuracy.
    pstPosterior = trnEstimatorTrainer.build_posterior(estEstimator)

SBI训练网络以及从后验分布采样时,会在当前工作目录建立sbi-logs目录,若有需要,可通过切换工作目录的方式,将该目录置入给定的临时目录中。:

    # Change working directory to avoid sbi-logs spamming
    import os
    import shutil
    import uuid
    sSBICacheDir = "temp/sbi/"
    sPreviousWorkDir = os.getcwd()
    sCurrentSBICacheDir = sSBICacheDir + str(uuid.uuid4()) + "/"
    os.makedirs(sCurrentSBICacheDir, exist_ok=True)
    os.chdir(sCurrentSBICacheDir)
 
    # Training, Posterior Building, and Sampling
    # ...
 
    # Revert working directory
    os.chdir(sPreviousWorkDir)
    shutil.rmtree(sCurrentSBICacheDir, ignore_errors=True)

调用build_posterior()时可以传入sample_with参数,决定使用何种算法进行采样。根据SBI的推荐

  • 使用NPE模型时:
  • 使用NLE或NRE模型时:
    • 若参数量小于3:使用RejectionPosteriorMCMCPosterior
    • 若参数量在3~10之间:使用MCMCPosterior
    • 若参数量大于10:使用VIPosterior
      • VIPosterior需要根据模型参数和仿真数据进行二次训练:
        • pstPosterior.train_amortized(arrParams, arrFeatures)
      • 可以进一步使用ImportanceSamplingPosterior优化采样准确率:
        • pstPosterior = sbi.inference.ImportanceSamplingPosterior(pstPosterior.potential_fn, dstSBIPriorDistribution, oversampling_factor=32)

根据观测值估计参数分布

后向推理器提供sample(sample_shape, x)函数,用于根据观测值x抽取具有给定个数的后验分布估计(Inferred Posterior Distribution):

    # Generate observed features
    arrInputParamsReal = [4.5, 40]
    arrObservedFeatures = torch.as_tensor(SBIParamToFeatures(arrInputParamsReal, True))
    print(f"Observed features: {arrInputParamsReal} -> {arrObservedFeatures}")
 
    # Estimating
    arrOptimizationParameters = pstPosterior.sample((100,), x=arrObservedFeatures).cpu().detach().numpy()
    arrOptimizationParameters = arrOptimizationParameters[0]
    print(f"Estimated parameter distribution: {arrOptimizationParameters}")
    print(f"Observation with estimated parameter distribution: {SBIParamToFeatures(arrOptimizationParameters, True)}")

使用sbi.analysis.pairplot()可视化推理结果:

    # Inspecting
    import sbi.analysis 
    import matplotlib.pyplot as plt
    arrOptimizationParameters = pstPosterior.sample((1000,), x=arrObservedFeatures).cpu().detach().numpy()
    sbi.analysis.pairplot(samples=arrOptimizationParameters,
        limits=[[1,10],[25,45]],
        ticks=[[1,10],[25,45]],
        figsize=(5, 5),
        points=[arrInputParamsReal],
        labels=["P0", "P1"])
    plt.show()
#End If

杂项

使用GPU

PyTorch支持使用GPU进行训练。只要在构建训练器时将device参数设为“cuda”或包含设备号的“cuda:0”等形式即可。同时需要通过sbi.utils.BoxUniform.to()方法将先验边界移动到目标设备上:

    # Detect availability of CUDA and GPU devices
    # To use SBI with GPU, see *https://sbi.readthedocs.io/en/latest/how_to_guide/07_gpu_training.html
    IsCUDAAvailable = torch.cuda.is_available()
    nGPUDevices = torch.cuda.device_count()
    iTensorSize = arrParams.element_size() * arrParams.numel()
    iTensorSize += arrFeatures.element_size() * arrFeatures.numel()
    sTrainingDevice = "cpu"
    if IsCUDAAvailable and (nGPUDevices > 0) :
        # Check GPU memory
        prpGPUproperties = torch.cuda.get_device_properties("cuda:0")
        iGPUMemory = prpGPUproperties.total_memory
        if iTensorSize < iGPUMemory * 1.5 :
            SharedUtil.LogInfo(f"CUDA is availabe on your system, and {nGPUDevices} GPU device(s) found")
            iGPUMemoryGB = round(iGPUMemory / 1024**3, 1)
            SharedUtil.LogInfo(f"Taining using CUDA device 0: {prpGPUproperties.name}, {iGPUMemoryGB} GB of memory")
            sTrainingDevice = "cuda:0"
        #End If
    #End If
    SharedUtil.LogInfo(f"SBI estimator will be trained on device {sTrainingDevice}")
 
    # Move prior distribution to the target device in-place
    dstSBIPriorDistribution.to(sTrainingDevice)
 
    # Train a neural network (neural density estimator, or neural classifier) to approximate p(Param | Observed), likelihood p(Observed | Param), or other variants
    # For estimating methods, see: https://sbi.readthedocs.io/en/latest/how_to_guide/06_choosing_inference_method.html and https://sbi.readthedocs.io/en/latest/api_reference/training.html
    #     NPE (Neural Posterior Estimation): Trains a deep neural density estimator to directly approximate the posterior p(Param | Observed).
    #         NPE is the only method which estimates the posterior directly. This means that it will not require further sampling steps (such as MCMC) after training. This makes NPE the fastest to sample from the posterior estimate after training. In addition, it can use an embedding network to learn summary statistics (just like NRE, but unlike NLE).
    #     NLE (Neural Likelihood Estimation): Trains a neural network to approximate the likelihood p(Observed | Param) using a conditional density estimator (normalizing flow).
    #         NLE learns the likelihood. This allows NLE to emulate the simulator after training. To sample from the posterior, NLE is combined with MCMC or variational inference, which makes it slower to sample. An advantage of NLE is that, if you have multiple independent identically distributed (IID) observations, then it can be much more simulation efficient than NPE
    #     NRE (Neural Ratio Estimation): Trains a neural classifier to estimate the likelihood-to-evidence ratio r(Param, Observed) = p(Observed | Param) / p(Observed) by distinguishing between samples from the joint distribution p(Param, Observed) and samples from the marginals p(Param)p(Observed). Posterior sampling is then performed via MCMC, rejection sampling, or variational inference using the estimated ratio.
    #         NRE learns the likelihood-to-evidence ratio. It trains only a classifier (unlike NPE and NLE, which train a generative model), which makes it fastest to train. Like NLE, NRE has to be combined MCMC or variational inference to sample. Like NPE, NRE can use an embedding network to learn summary statistics.
    # SBI explicit recommendations:
    #     If have high-dimensional simulation-outputs (images, time-series,…), use NPE or NRE.
    #     If you want fast sampling, use NPE.
    #     If you have IID observations across many trials (and simulations are relatively expensive), use NRE or NLE.
    #     If you want to perform inference for many observations, use amortized methods (which is the default of the sbi package).
    #     If you want to perform inference for few observations (<10), then start with amortized methods and switch to sequential methods only if the simulation cost is too high.
    #     Do not use SNPE_A or SNPE_B unless you are an expert user and know what you are doing.
    trnEstimatorTrainer = sbi.inference.NPE_C(prior=dstSBIPriorDistribution, device=sTrainingDevice)
    trnEstimatorTrainer = trnEstimatorTrainer.append_simulations(arrParams, arrFeatures)
    estEstimator = trnEstimatorTrainer.train()
 
    # Build posterior p(Param | Observed) sampler from the trained neural network
    # For sampler selection, see: *https://sbi.readthedocs.io/en/latest/how_to_guide/09_sampler_interface.html
    #     MCMCPosterior: very accurate, but can be slow if you have many parameters.
    #     VIPosterior: can be much faster if you have many parameter. May be inaccurate. You may need to train VIPosterior before sampling.
    #     RejectionPosterior: accurate and fast, but only for very few parameters (typically less than 3).
    #     ImportanceSamplingPosterior: typically inaccurate, but can be very useful to improve the accuracy of a VIPosterior (see above).
    #     For NPE:
    #         use the DirectPosterior. If you have a likelihood available, you can refine the posterior with importance sampling
    #     For NLE or NRE:
    #         If you have very few parameters (<3), use RejectionPosterior or MCMCPosterior
    #         If you have a medium number of parameters (3-10), use MCMCPosterior
    #         If you have many parameters (>10) and the MCMCPosterior is too slow, use the VIPosterior. Optionally combine the VIPosterior with an ImportanceSamplingPosterior to improve its accuracy.
    pstPosterior = trnEstimatorTrainer.build_posterior(estEstimator)
 
    # Generate observed features
    arrInputParamsReal = [4.5, 40]
    arrObservedFeatures = torch.as_tensor(SBIParamToFeatures(arrInputParamsReal, True)).to(torch.device(sTrainingDevice))
    print(f"Observed features: {arrInputParamsReal} -> {arrObservedFeatures}")
 
    # Estimating
    arrOptimizationParameters = pstPosterior.sample((100,), x=arrObservedFeatures).cpu().detach().numpy()

常数问题

如果模型存在上界与下界相同的参数(常数参数),SBI可能在执行build_posterior()等函数时失败,并产生错误:

AssertionError: Original and re-transformed parameters must be close to each other.

这是因为SBI的变换函数涉及Z-Score的变换与反变换。对于常数(Z-Score为0)的情况,反变换时出现“0/0”型错误,返回了“nan”。

建议在进入SBI流程前,将参数数组分割为可变部分与常数部分,在SBI推理结束后再合并。

临时处理(可能导致结果异常)时,可以编辑torch/distributions/transforms.py文件中AffineTransform类的_inverse()函数,实现一个简单的安全除法(将naninf置为0):

    def _inverse(self, y):
        # Use safe div
        # See *https://duoke360.com/post/35074
        inv_scale = 1. / self.scale
        inv_scale = torch.nan_to_num(inv_scale, nan=0, posinf=245.0, neginf=-245.0)
        return (y - self.loc) * inv_scale

同时编辑torch/distributions/uniform.py文件中Uniform类的log_prob()函数,实现一个简单的安全除法:

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        lb = self.low.le(value).type_as(self.low)
        ub = self.high.gt(value).type_as(self.low)
        result = torch.log(lb.mul(ub)) - torch.log(self.high - self.low)
        return torch.nan_to_num(result, nan=0, posinf=245.0, neginf=-245.0)

参考资料

https://sbi.readthedocs.io/en/latest/tutorials/Example_00_HodgkinHuxleyModel.html

https://sbi.readthedocs.io/en/latest/how_to_guide/06_choosing_inference_method.html

https://sbi.readthedocs.io/en/latest/how_to_guide/07_gpu_training.html

https://sbi.readthedocs.io/en/latest/how_to_guide/09_sampler_interface.html

bolfazl Ziaeemehr, Marmaduke Woodman, Lia Domide, Spase Petkoski, Viktor Jirsa, Meysam Hashemi (2025) Virtual Brain Inference (VBI), a flexible and integrative toolkit for efficient probabilistic inference on whole-brain models eLife 14:RP106194 https://doi.org/10.7554/eLife.106194.4

Meysam Hashemi et al 2024 Mach. Learn.: Sci. Technol. 5 035019. https://doi.org/10.1088/2632-2153/ad6230

Meysam Hashemi, Anirudh N. Vattikonda, Jayant Jha, Viktor Sip, Marmaduke M. Woodman, Fabrice Bartolomei, Viktor K. Jirsa. Amortized Bayesian inference on generative dynamical network models of epilepsy using deep neural density estimators. Neural Networks, Volume 163, 2023, https://doi.org/10.1016/j.neunet.2023.03.040

it
除非特别注明,本页内容采用以下授权方式: Creative Commons Attribution-ShareAlike 3.0 License