如何修正神经代理（Neural Surrogates）中BatchNormalization维度不匹配错误

问题背景 #

构建带BatchNormalization的神经代理模型时出现维度不匹配错误，调整Dense层和BatchNorm层维度后仍存在其他问题，需解决该报错。

训练数据 #

x_train = 
100-element Vector{Tuple{Float64, Float64}}:
 (0.1875, 3.1875)
 (4.1875, 7.1875)
 (6.1875, 1.1875)
 (2.1875, 5.1875)
 (3.1875, 0.1875)
 (7.1875, 4.1875)
 (5.1875, 2.1875)
 (1.1875, 6.1875)
 (1.6875, 1.6875)
 (5.6875, 5.6875)
 ⋮
 (3.34375, 5.90625)
 (2.34375, 0.90625)
 (6.34375, 4.90625)
 (4.34375, 2.90625)
 (0.34375, 6.90625)
 (0.46875, 0.53125)
 (4.46875, 4.53125)
 (6.46875, 2.53125)
 (2.46875, 6.53125)

y_train = 
100-element Vector{Float64}:
 10.1953125
 69.1953125
 39.6953125
 31.6953125
 10.1953125
 69.1953125
 31.6953125
 39.6953125
  5.6953125
 64.6953125
 ⋮
 46.064453125
  6.314453125
 64.314453125
 27.314453125
 47.814453125
  0.501953125
 40.501953125
 48.251953125
 48.751953125

原模型代码 #

# Define the neural network model with normalization
model1 = Chain(
    Dense(2, 100),
    BatchNorm(10, relu),
    Dense(10, 10),
    BatchNorm(10),
    Dense(10, 1)
)

neural_surrogate = NeuralSurrogate(
    x_train,
    y_train,
    lower_bound,
    upper_bound,
    model = model1,
    n_echos = 10
)

错误信息 #

layer BatchNorm(10, relu) expects ndims(input) > 1, but got 100-element Vector{Float32}

问题分析与解决方案 #

错误源于两个核心问题：

1. 层维度不匹配：第一个Dense层输出维度为100，但后续BatchNorm层指定维度为10，两者不匹配导致输出张量维度异常，触发BatchNorm对输入维度的要求报错。
2. 输入数据格式错误：x_train是Tuple类型的向量，Flux神经网络默认要求输入为特征数×样本数的矩阵格式，而非Tuple向量，否则会被解析为1维张量，不符合BatchNorm对输入至少2维的要求。

修正步骤： #

1. 修正模型层维度匹配

确保每个BatchNorm层的维度与前一层输出维度一致，同时Dense层输入维度匹配前一层输出：

model1 = Chain(
    Dense(2, 100),
    BatchNorm(100, relu),  # 与前一层Dense输出维度100匹配
    Dense(100, 10),        # 输入维度匹配前一层BatchNorm输出的100
    BatchNorm(10),
    Dense(10, 1)
)

2. 转换训练数据格式

将x_train转换为Flux要求的矩阵格式，并统一数据类型为Float32（与模型默认权重类型一致）：

# 将Tuple向量转换为2×100的矩阵，每列对应一个样本
x_train_matrix = hcat([collect(x) for x in x_train]...) |> Float32
# 将y_train转换为Float32类型
y_train_vector = y_train |> Float32

3. 创建修正后的神经代理模型

使用处理后的训练数据初始化模型：

neural_surrogate = NeuralSurrogate(
    x_train_matrix,
    y_train_vector,
    lower_bound,
    upper_bound,
    model = model1,
    n_echos = 10
)

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