Flux.jl自定义优化器实现难题：RSO无梯度优化器训练CNN时的单权重更新与损失评估

我明白你在实现RSO（基于采样的无梯度优化器）时遇到的核心痛点：Flux.jl的Zygote.Params是不可变的包装结构，没法直接修改单个权重来测试W+ΔWj、当前W、W-ΔWj三种状态下的损失。下面给你两个实用的解决方案，完全适配Flux的参数体系，同时修正你当前代码里的一些问题：

方案1：参数扁平化（推荐，更直观） #

Flux提供了Flux.destructure和restructure工具，可以把模型的所有参数转换成一维数组，修改后再重构回模型。这种方式不用关心参数的层级结构，操作单个权重非常方便。

实现步骤： #

1. 先将模型参数扁平化，得到一维数组和重构函数
2. 为每个权重元素建立从模型参数到一维数组的索引映射（提前遍历一次即可）
3. 对目标权重生成三种参数版本，分别计算损失
4. 选择损失最小的参数版本，重构回模型

代码示例（整合到你的RSO函数中）： #

function RSO(X, L, C, model, batch_size, device)
    # 数据归一化修正：用训练集均值/标准差，而非全局总和
    X_mean = mean(X)
    X_std = std(X)
    X .= (X .- X_mean) ./ X_std

    train_loader = DataLoader((X, L), batchsize=batch_size, shuffle=true)

    # 1. 扁平化模型参数，获取重构函数
    flat_params, re = Flux.destructure(model)
    flat_params = flat_params |> device

    # 2. 建立参数到一维数组的索引映射：记录每个参数在flat_params中的起始/结束索引
    param_indices = []
    current_idx = 1
    for param in Flux.params(model)
        param_len = length(param)
        push!(param_indices, (start=current_idx, stop=current_idx + param_len - 1))
        current_idx += param_len
    end

    # 3. 初始化权重标准差（按论文要求，每个层单独计算）
    σ_d = []
    for (i, param) in enumerate(Flux.params(model))
        # 卷积层fan_in = 输入通道数 * 核宽 * 核高
        fan_in = size(param)[1] * size(param)[2] * size(param)[3]
        init_std = sqrt(2 / fan_in)
        param.data .= randn!(param.data) .* init_std
        push!(σ_d, init_std) # 用初始化标准差作为ΔWj的基准
    end

    # 4. RSO权重更新循环
    for _ in 1:C
        # 反向遍历每一层（按论文要求从最后一层到第一层）
        for d in reverse(1:length(Flux.params(model)))
            param = Flux.params(model)[d]
            idx_range = param_indices[d]
            param_std = σ_d[d]

            # 遍历当前参数的每个元素
            for j in 1:length(param)
                # 随机采样一个mini-batch
                batch_idx = rand(1:length(train_loader))
                x, l = train_loader[batch_idx]
                x = x |> device
                l = l |> device

                # 计算当前元素在flat_params中的索引
                flat_idx = idx_range.start + j - 1

                # 生成三种参数版本
                ΔWj = randn(device) * param_std
                # W+ΔWj
                flat_plus = copy(flat_params)
                flat_plus[flat_idx] += ΔWj
                model_plus = re(flat_plus) |> device
                loss_plus = logitcrossentropy(model_plus(x), l, agg=sum)

                # 当前W
                loss_current = logitcrossentropy(model(x), l, agg=sum)

                # W-ΔWj
                flat_minus = copy(flat_params)
                flat_minus[flat_idx] -= ΔWj
                model_minus = re(flat_minus) |> device
                loss_minus = logitcrossentropy(model_minus(x), l, agg=sum)

                # 选择损失最小的参数版本
                min_loss, min_idx = findmin([loss_plus, loss_current, loss_minus])
                if min_idx == 1
                    flat_params = flat_plus
                elseif min_idx == 3
                    flat_params = flat_minus
                end

                # 重构回模型
                model = re(flat_params) |> device
            end
        end
    end

    return model
end

方案2：直接操作参数底层数组（更高效，适合大模型） #

Flux的Param对象其实是对底层数组的包装，你可以通过.data直接访问和修改数组元素（GPU上的数组也可以直接操作）。这种方式不用频繁重构模型，效率更高。

代码示例： #

function RSO(X, L, C, model, batch_size, device)
    # 数据归一化修正
    X_mean = mean(X)
    X_std = std(X)
    X .= (X .- X_mean) ./ X_std

    train_loader = DataLoader((X, L), batchsize=batch_size, shuffle=true)

    # 初始化权重和标准差
    σ_d = []
    for param in Flux.params(model)
        fan_in = size(param)[1] * size(param)[2] * size(param)[3]
        init_std = sqrt(2 / fan_in)
        param.data .= randn!(param.data) .* init_std
        push!(σ_d, init_std)
    end

    # RSO更新循环
    for _ in 1:C
        # 反向遍历每一层
        for (d, param) in enumerate(reverse(Flux.params(model)))
            param_std = σ_d[end - d + 1] # 对应反向层的标准差

            # 遍历参数的每个元素（用线性索引）
            for j in 1:length(param)
                # 采样mini-batch
                x, l = train_loader[rand(1:length(train_loader))]
                x = x |> device
                l = l |> device

                original_val = param.data[j]
                ΔWj = randn(device) * param_std

                # 测试W+ΔWj
                param.data[j] = original_val + ΔWj
                loss_plus = logitcrossentropy(model(x), l, agg=sum)

                # 测试当前W
                param.data[j] = original_val
                loss_current = logitcrossentropy(model(x), l, agg=sum)

                # 测试W-ΔWj
                param.data[j] = original_val - ΔWj
                loss_minus = logitcrossentropy(model(x), l, agg=sum)

                # 恢复最优值
                min_loss, min_idx = findmin([loss_plus, loss_current, loss_minus])
                if min_idx == 1
                    param.data[j] = original_val + ΔWj
                elseif min_idx == 2
                    param.data[j] = original_val
                else
                    param.data[j] = original_val - ΔWj
                end
            end
        end
    end

    return model
end

你当前代码的几个关键修正点 #

• 权重初始化错误：之前的id = wj只是给局部变量赋值，没有修改模型参数，应该直接操作param.data
• 归一化错误：用sum(X)计算均值是错误的，需替换为mean(X)
• ΔWj生成逻辑：论文中ΔWj的标准差应与对应层的权重初始化标准差关联，而非所有层权重的全局标准差
• 损失选择逻辑：argmin(F(x,l,W+ΔWj), ...)写法不成立，需先计算三个损失值，再选择对应最小损失的参数状态

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