如何在Optuna中从加权分类分布采样参数？

针对你需要给分类搜索项设置权重、引导采样器优先探索大子空间的需求，以下是两种可落地的实现方案：

方案1：自定义带加权先验的TPESampler（最推荐） #

通过继承原生TPESampler扩展加权逻辑，既保留分类参数的原生结构，也不会破坏TPE的自适应优化能力：

import optuna
from optuna.samplers import TPESampler
from typing import Optional, Dict, Any

class WeightedCategoricalTPESampler(TPESampler):
    def __init__(self, categorical_weights: Optional[Dict[str, Dict[Any, float]]] = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 权重配置格式：{参数名: {分类选项: 权重值}}，权重越高采样优先级越高
        self.categorical_weights = categorical_weights or {}

    def sample_independent(self, study, trial, param_name, param_distribution):
        # 仅对指定了权重的分类参数应用自定义采样逻辑
        if param_name in self.categorical_weights and isinstance(param_distribution, optuna.distributions.CategoricalDistribution):
            choices = param_distribution.choices
            weights = [self.categorical_weights[param_name].get(c, 1.0) for c in choices]
            # 前N次试验按自定义权重采样，N可根据分类总数调整，推荐设为分类数的1~2倍
            if trial.number < len(choices) * 2:
                import random
                return random.choices(choices, weights=weights, k=1)[0]
        # 其余场景走原生TPE采样逻辑，自动根据试验结果调整采样概率
        return super().sample_independent(study, trial, param_name, param_distribution)

使用方式：

# 示例：给nn设置100倍权重，引导采样器优先探索神经网络子空间
sampler = WeightedCategoricalTPESampler(
    categorical_weights={
        "network_type": {"linear": 1, "svm": 1, "nn": 100}
    }
)
study = optuna.create_study(sampler=sampler)

优势：

• 分类参数结构完整保留，后续结果分析无需额外映射处理
• 采样器可以正常识别分类空间，避免了浮点数映射方案中「数值接近的浮点数对应完全不同分类」的逻辑缺陷
• 前期按人工先验权重采样，保证大子空间获得足够的探索资源，达到设定步数后自动切换为原生TPE的性能驱动采样逻辑，自动提高表现更优的分类项的采样概率

方案2：目标函数内加采样限制（无需修改采样器） #

如果仅需要临时限制小子空间的采样次数，不需要灵活调整权重，可以直接在目标函数中加剪枝逻辑：

def objective(trial):
    network_type = trial.suggest_categorical("network_type", ["linear", "svm", "nn"])
    # 前100次试验限制linear和svm最多各采样5次，避免占用过多资源
    if trial.number < 100:
        completed_trials = study.get_trials(states=[optuna.trial.TrialState.COMPLETE])
        linear_cnt = sum(1 for t in completed_trials if t.params.get("network_type") == "linear")
        svm_cnt = sum(1 for t in completed_trials if t.params.get("network_type") == "svm")
        if (network_type == "linear" and linear_cnt >=5) or (network_type == "svm" and svm_cnt >=5):
            raise optuna.TrialPruned()
    # 后续正常执行模型训练与评估逻辑
    ...

千级分类场景优化建议 #

• 初始化TPESampler时可以调大prior_weight参数（默认值为1.0，推荐调整到10~100），让TPE更尊重你设置的初始先验，不会过快切换到纯性能驱动的采样逻辑
• 可以先对上千个分类项做分层聚合，先按组做加权采样，再在组内做细粒度采样，进一步降低搜索复杂度

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