BoTorch中qKnowledgeGradient与SaasGP适配遇形状不匹配错误

错误现象 #

尝试将qKnowledgeGradient与SaasFullyBayesianSingleTaskGP结合，通过继承FantasizeMixin自定义模型后，运行时触发形状不匹配错误：

RuntimeError: shape '[2, 1, 16, 1]' is invalid for input of size 64

后续修改num_fantasies为64并重写condition_on_observations后，出现新错误：

Output shape not equal to that of weights. Output shape is 1 and weights are torch.Size([64])

核心成因 #

1. SAAS GP的批次维度未被处理：训练后的SaasFullyBayesianSingleTaskGP带有对应MCMC样本的批次维度（本例中thinning=16，256个样本后批次维度为(16,)），而FantasizeMixin默认实现未适配多批次模型，导致幻想采样的张量形状与模型批次维度冲突。
2. 后验变换权重形状错误：原代码中weights = torch.ones(2)，但模型是单输出，权重形状必须与输出维度一致（即(1,)）。
3. 自定义condition_on_observations逻辑混乱：错误地调整了维度顺序和展开逻辑，导致张量形状进一步错位。

正确实现方案 #

修正带幻想功能的SAAS GP类 #

重写fantasize和condition_on_observations，确保MCMC批次维度与幻想维度正确对齐：

class SaasFullyBayesianSingleTaskGPWithFantasy(SaasFullyBayesianSingleTaskGP, FantasizeMixin):
    def fantasize(
        self,
        X: torch.Tensor,
        sampler: Optional[MCSampler] = None,
        num_fantasies: int = 2,
        **kwargs,
    ) -> Model:
        if sampler is None:
            # 采样形状需包含模型批次维度 + 幻想维度
            sampler = SobolQMCNormalSampler(
                sample_shape=self.batch_shape + torch.Size([num_fantasies]),
                collapse_batch_dims=False,
            )
        X = torch.as_tensor(
            X, dtype=self.train_inputs[0].dtype, device=self.train_inputs[0].device
        )
        # 调用FantasizeMixin的fantasize，传递正确的采样器
        fantasy_model = FantasizeMixin.fantasize(self, X, sampler=sampler, **kwargs)
        # 保持幻想模型的批次维度与原模型一致
        fantasy_model.batch_shape = self.batch_shape + torch.Size([num_fantasies])
        return fantasy_model

    def condition_on_observations(self, X: torch.Tensor, Y: torch.Tensor, **kwargs) -> Model:
        # 处理X：确保X的批次维度与模型匹配，添加幻想维度
        if X.ndim == 2:
            # 扩展X到模型批次维度 + 幻想维度
            X = X.expand(self.batch_shape + X.shape)
        # 处理Y：确保Y的形状与X对齐
        if Y.ndim == 2:
            Y = Y.expand(self.batch_shape + Y.shape)
        # 调用父类方法，注意SAAS GP的condition_on_observations需要正确的批次输入
        return super().condition_on_observations(X, Y, **kwargs)

修正后的运行代码 #

调整后验变换的权重形状，确保与模型单输出匹配：

import numpy as np
import math
import torch
from botorch.models import SaasFullyBayesianSingleTaskGP, FantasizeMixin
from botorch.fit import fit_fully_bayesian_model_nuts
from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
from botorch.transforms import ScalarizedPosteriorTransform
from botorch.optim import optimize_acqf
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.utils.sampling import SobolEngine
from typing import Optional

# 嵌入100维的Branin函数
lb = np.hstack((-5 * np.ones(50), 0 * np.ones(50)))
ub = np.hstack((10 * np.ones(50), 15 * np.ones(50)))

def branin100(x):
    assert (x <= ub).all() and (x >= lb).all()
    x1, x2 = x[19], x[64]
    t1 = x2 - 5.1 / (4 * math.pi ** 2) * x1 ** 2 + 5 / math.pi * x1 - 6
    t2 = 10 * (1 - 1 / (8 * math.pi)) * np.cos(x1)
    return t1 ** 2 + t2 + 10

def run_saasbo_botorch():
    torch.manual_seed(0)
    dtype = torch.double
    device = "cpu"
    dim = 100
    lb_torch = torch.zeros(dim, dtype=dtype)
    ub_torch = torch.ones(dim, dtype=dtype)
    bounds = torch.stack([lb_torch, ub_torch])

    def f(x): return branin100(x)

    # Initial Sobol samples
    sobol = SobolEngine(dim, scramble=True, seed=0)
    X = sobol.draw(4).to(dtype=dtype)  # 4 initial points
    Y = torch.tensor(
        [f(lb + (ub - lb) * x.cpu().numpy()) for x in X],
        dtype=dtype
    ).unsqueeze(-1)

    train_Y = (Y - Y.mean()) / Y.std()

    # Fit SAAS GP
    model = SaasFullyBayesianSingleTaskGPWithFantasy(X, train_Y)
    fit_fully_bayesian_model_nuts(
        model, warmup_steps=512, num_samples=256, thinning=16
    )

    # 修正：权重形状匹配单输出模型
    weights = torch.ones(1, dtype=dtype)
    post_tf = ScalarizedPosteriorTransform(weights=weights)

    # Define KG acquisition
    qkg = qKnowledgeGradient(
        model=model,
        num_fantasies=2,
        current_value=train_Y.min(),
        posterior_transform=post_tf,
    )

    # Optimize acquisition
    candidate, _ = optimize_acqf(
        acq_function=qkg,
        bounds=bounds,
        q=1,
        raw_samples=1,
        num_restarts=1,
    )
    print("候选点:", candidate)

run_saasbo_botorch()

关键说明 #

• 采样器的sample_shape必须包含模型的batch_shape（MCMC样本批次）和num_fantasies，确保幻想采样的张量与模型批次维度对齐。
• 后验变换的权重形状必须与模型输出维度一致（单输出用(1,)，多输出对应输出维度）。
• condition_on_observations方法需要确保输入X/Y的批次维度与模型的批次+幻想维度匹配，避免形状冲突。

内容的提问来源于stack exchange，提问作者Helena