使用自定义Lux层运行DiffEqFlux的NeuralODE时遇MethodError

问题根源 #

错误提示已经点明核心：你定义的LuxNeurLayer只实现了接受AbstractMatrix类型输入的调用方法，但NeuralODE在计算过程中传入的是Vector{Float32}（也就是你的u0），导致方法不匹配触发报错。

修复步骤 #

1. 让层同时支持向量和矩阵输入 #

修改自定义层的调用方法，使其兼容AbstractVecOrMat（同时覆盖向量和矩阵类型），统一处理逻辑后再根据输入类型返回对应格式的结果：

function (l::LuxNeurLayer)(x::AbstractVecOrMat, ps, st::NamedTuple)
    # 把向量转成列矩阵，统一后续计算逻辑
    x_mat = reshape(x, :, size(x, ndims(x)))
    y = similar(x_mat, l.out_dims, size(x_mat, 2))
    num_neurs = Int(l.out_dims / 2)  # 补上之前未定义的num_neurs变量
    for batch_idx in 1:size(x_mat, 2)
        x_col = x_mat[:, batch_idx]
        y_col = zeros(l.out_dims)
        for i=1:num_neurs
            connections = zeros(l.in_dims)
            for j=1:l.in_dims
                # 这里注意：原代码x[2*j]会超出u0的维度，后续要修正维度匹配问题
                connections[j] = ps.gain[i,j] / (1 + exp(-0.4 * (x_col[j] + 0.6) ))
            end
            y_col[2*i-1] = -y_col[2*i-1] + 0.5*tanh(y_col[2*i-1]) + sum(connections)
            y_col[2*i]   = (y_col[2*i-1] - y_col[2*i]) / 5.0
        end
        y[:, batch_idx] = y_col
    end
    # 如果输入是向量，输出转回向量；矩阵则保持原格式
    return ndims(x) == 1 ? vec(y) : y, st
end

2. 修正维度不匹配问题 #

你的u0是2维向量，但自定义层的in_dims=20，这会导致索引访问越界，必须二选一调整：

• 要么把u0改成20维：

  u0 = rand(Float32, 20)
  

• 要么调整层的输入维度为2，匹配u0：

  l = LuxNeurLayer(2,4)
  

3. 补充缺失的时间参数 #

原代码里没定义tspan和tsteps，NeuralODE需要这两个参数：

tspan = (0.0f0, 1.0f0)
tsteps = 0.0f0:0.1f0:1.0f0

完整修复后的代码 #

using Lux, DiffEqFlux, DifferentialEquations, Optimization, OptimizationOptimJL, Random, Plots, ComponentArrays

struct LuxNeurLayer{F1} <: Lux.AbstractExplicitLayer
    in_dims::Int
    out_dims::Int
    init_gain::F1
end

function LuxNeurLayer(in_dims::Int, out_dims::Int; init_gain=Lux.glorot_uniform)
    return LuxNeurLayer{typeof(init_gain)}(in_dims, out_dims, init_gain)
end

# 调整in_dims为2，匹配u0的维度
l = LuxNeurLayer(2,4)

function Lux.initialparameters(rng::AbstractRNG, l::LuxNeurLayer)
    return (gain=l.init_gain(rng, Int(l.out_dims/2), l.in_dims),)
end

Lux.initialstates(::AbstractRNG, ::LuxNeurLayer) = NamedTuple()

function (l::LuxNeurLayer)(x::AbstractVecOrMat, ps, st::NamedTuple)
    x_mat = reshape(x, :, size(x, ndims(x)))
    y = similar(x_mat, l.out_dims, size(x_mat, 2))
    num_neurs = Int(l.out_dims / 2)
    for batch_idx in 1:size(x_mat, 2)
        x_col = x_mat[:, batch_idx]
        y_col = zeros(l.out_dims)
        for i=1:num_neurs
            connections = zeros(l.in_dims)
            for j=1:l.in_dims
                # 修正索引，避免越界
                connections[j] = ps.gain[i,j] / (1 + exp(-0.4 * (x_col[j] + 0.6) ))
            end
            y_col[2*i-1] = -y_col[2*i-1] + 0.5*tanh(y_col[2*i-1]) + sum(connections)
            y_col[2*i]   = (y_col[2*i-1] - y_col[2*i]) / 5.0
        end
        y[:, batch_idx] = y_col
    end
    return ndims(x) == 1 ? vec(y) : y, st
end

rng = Random.default_rng()
Random.seed!(rng, 0)

p, st = Lux.setup(rng, l)
pinit = ComponentArray(p)

# 补充时间参数
tspan = (0.0f0, 1.0f0)
tsteps = 0.0f0:0.1f0:1.0f0

prob_neuralode = NeuralODE(l, tspan, Tsit5(), saveat = tsteps)

u0 = Float32[2.0; 0.0]
# 现在可正常执行
result = prob_neuralode(u0, pinit, st)
println(result)

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