如何在Python中借助鲸鱼优化算法筛选最优优化器

核心逻辑是将4种优化器的选择映射为WOA的离散搜索维度，搭配优化器核心超参数作为连续搜索维度，以模型在验证集上的性能作为适应度，通过WOA的全局迭代搜索得到最优配置。

实现流程 #

• 定义搜索空间：离散维度对应4种优化器编号（0=随机梯度下降/SGD、1=ADAM、2=RMSprop、3=AdaGrad），连续维度对应学习率、权重衰减等通用超参数的取值范围
• 适配WOA编码规则：对WOA输出的连续位置值做离散映射，优化器编号位通过四舍五入取整+边界裁剪得到合法值，超参数位直接裁剪到预设取值区间
• 定义适应度函数：将位置值解码为具体优化器配置，初始化模型后用对应优化器训练固定轮次，返回验证集损失作为适应度（取值越小性能越好）
• 运行WOA迭代流程，每轮更新全局最优适应度对应的配置，迭代结束后输出最优优化器类型及对应超参数

可直接运行的Python代码 #

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import SGD, Adam, RMSprop, Adagrad

# 基础配置
SEARCH_SPACE = {
    "optimizer_type": [0, 1, 2, 3],
    "lr": (1e-5, 1e-1),
    "weight_decay": (1e-6, 1e-3)
}
WOA_POPULATION = 10  # 鲸鱼种群规模
WOA_ITER = 20  # 迭代轮次
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 测试模型（实际使用时替换为自己任务的模型即可）
class TaskModel(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=20, output_dim=2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    def forward(self, x):
        return self.net(x)

# WOA优化器选择器核心实现
class WOA_OptSelector:
    def __init__(self, train_loader, val_loader, loss_fn):
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.loss_fn = loss_fn
        self.dim = 3  # 搜索维度：优化器类型、学习率、权重衰减
        self.lb = np.array([0, SEARCH_SPACE["lr"][0], SEARCH_SPACE["weight_decay"][0]])
        self.ub = np.array([3, SEARCH_SPACE["lr"][1], SEARCH_SPACE["weight_decay"][1]])
        self.best_pos = None
        self.best_fitness = float("inf")

    def decode_pos(self, pos):
        # 连续位置转实际可用配置
        opt_type = int(np.clip(np.round(pos[0]), 0, 3))
        lr = np.clip(pos[1], self.lb[1], self.ub[1])
        wd = np.clip(pos[2], self.lb[2], self.ub[2])
        return opt_type, lr, wd

    def init_pop(self):
        # 初始化鲸鱼种群
        pop = np.random.uniform(0, 1, size=(WOA_POPULATION, self.dim))
        for i in range(self.dim):
            pop[:, i] = self.lb[i] + pop[:, i] * (self.ub[i] - self.lb[i])
        return pop

    def calc_fitness(self, pos):
        opt_type, lr, wd = self.decode_pos(pos)
        model = TaskModel().to(DEVICE)
        # 匹配对应优化器
        if opt_type == 0:
            optimizer = SGD(model.parameters(), lr=lr, weight_decay=wd)
        elif opt_type == 1:
            optimizer = Adam(model.parameters(), lr=lr, weight_decay=wd)
        elif opt_type == 2:
            optimizer = RMSprop(model.parameters(), lr=lr, weight_decay=wd)
        else:
            optimizer = Adagrad(model.parameters(), lr=lr, weight_decay=wd)
        # 固定轮次训练评估
        model.train()
        for _ in range(3):
            for x, y in self.train_loader:
                x, y = x.to(DEVICE), y.to(DEVICE)
                optimizer.zero_grad()
                loss = self.loss_fn(model(x), y)
                loss.backward()
                optimizer.step()
        # 计算验证集损失作为适应度
        model.eval()
        total_loss, sample_cnt = 0, 0
        with torch.no_grad():
            for x, y in self.val_loader:
                x, y = x.to(DEVICE), y.to(DEVICE)
                batch_loss = self.loss_fn(model(x), y).item()
                total_loss += batch_loss * x.shape[0]
                sample_cnt += x.shape[0]
        return total_loss / sample_cnt

    def run(self):
        pop = self.init_pop()
        # 初始化全局最优
        for i in range(WOA_POPULATION):
            fit = self.calc_fitness(pop[i])
            if fit < self.best_fitness:
                self.best_fitness = fit
                self.best_pos = pop[i].copy()
        # 迭代搜索
        for t in range(WOA_ITER):
            a = 2 - t * (2 / WOA_ITER)
            a2 = -1 + t * (-1 / WOA_ITER)
            for i in range(WOA_POPULATION):
                r1, r2 = np.random.rand(), np.random.rand()
                A = 2 * a * r1 - a
                C = 2 * r2
                b, l = 1, (a2 - 1) * np.random.rand() + 1
                p = np.random.rand()
                for j in range(self.dim):
                    if p < 0.5:
                        if abs(A) >= 1:
                            rand_idx = np.random.randint(0, WOA_POPULATION)
                            D = abs(C * pop[rand_idx, j] - pop[i, j])
                            pop[i, j] = pop[rand_idx, j] - A * D
                        else:
                            D = abs(C * self.best_pos[j] - pop[i, j])
                            pop[i, j] = self.best_pos[j] - A * D
                    else:
                        D_leader = abs(self.best_pos[j] - pop[i, j])
                        pop[i, j] = D_leader * np.exp(b*l) * np.cos(2*np.pi*l) + self.best_pos[j]
                pop[i] = np.clip(pop[i], self.lb, self.ub)
                fit = self.calc_fitness(pop[i])
                if fit < self.best_fitness:
                    self.best_fitness = fit
                    self.best_pos = pop[i].copy()
            print(f"迭代轮次 {t+1}/{WOA_ITER}，当前最优验证损失：{self.best_fitness:.4f}")
        # 输出最终结果
        opt_type, best_lr, best_wd = self.decode_pos(self.best_pos)
        opt_name_map = {0:"SGD（随机梯度下降）", 1:"ADAM", 2:"RMSprop", 3:"AdaGrad"}
        print("="*50)
        print(f"最优优化器：{opt_name_map[opt_type]}")
        print(f"对应超参数：学习率={best_lr:.6f}，权重衰减={best_wd:.6f}")
        print(f"最优验证损失：{self.best_fitness:.4f}")
        return opt_name_map[opt_type], best_lr, best_wd

# 调用示例（实际使用时替换为自己的数据集即可）
if __name__ == "__main__":
    train_x = torch.randn(1000, 10)
    train_y = torch.randint(0, 2, (1000,))
    val_x = torch.randn(200, 10)
    val_y = torch.randint(0, 2, (200,))
    train_loader = DataLoader(list(zip(train_x, train_y)), batch_size=32, shuffle=True)
    val_loader = DataLoader(list(zip(val_x, val_y)), batch_size=32, shuffle=False)
    loss_fn = nn.CrossEntropyLoss()
    selector = WOA_OptSelector(train_loader, val_loader, loss_fn)
    selector.run()

使用说明 #

• 代码中的测试模型、模拟数据集、3轮训练评估都是占位逻辑，实际使用时替换为自己任务对应的模型、数据集即可，评估轮次建议设为任务常规训练轮次的1/5~1/3，平衡搜索速度和评估准确性
• 如果需要搜索优化器专属超参数（比如SGD的动量、Adam的beta值），直接扩展搜索空间维度，在位置解码、优化器初始化部分补充对应逻辑即可
• 种群规模和迭代轮次可根据搜索空间大小调整，搜索维度越多，对应的种群数、迭代数需要适当增大
• 如果任务用准确率、AUC这类越大越好的指标，把适应度返回值改为负指标值即可，保证WOA始终朝最小化适应度的方向搜索

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