基于Evolution Strategies的MNIST神经网络始终输出相同准确率问题（PyTorch）

核心问题诊断 #

模型始终收敛到随机猜测准确率（~0.098）且训练缓慢，核心由以下关键错误导致：

1. 权重初始化完全失效 #

m.weight.data.normal_(0.0, 0.00)

标准差设为0会让所有线性层权重初始化为全0，模型从一开始就无法学习任何特征，输出完全一致，只能随机猜测类别。

2. 最后一层激活函数错误 #

模型最后一层使用ReLU：

nn.Linear(50, 10),
nn.ReLU(),

CrossEntropyLoss要求输入是未经过激活的logits，ReLU会截断负logits为0，破坏分类任务所需的输出分布，导致模型无法区分类别。

3. 输入数据未归一化 #

MNIST像素值范围是0-255，直接输入会导致线性层输出值过大，ReLU提前饱和，ES优化器无法捕捉有效更新信号。

4. 损失函数数值不稳定 #

用1/loss_func作为fitness，当loss趋近于0时会导致fitness暴增，引发数值波动。更合理的方式是用负损失作为fitness（最大化负损失等价于最小化损失）。

5. 训练效率低下 #

每次计算fitness都用全量训练集（49000样本），50个个体需要50次全量前向传播，是训练缓慢的核心原因。

修复后的完整代码 #

# imports
import torch
import torch.nn as nn
from tqdm.notebook import tqdm
import numpy as np
from sklearn.model_selection import train_test_split
from keras.datasets import mnist

# 加载并预处理数据：归一化+划分数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x = np.concatenate((x_train, x_test)) / 255.0  # 归一化到0-1
y = np.concatenate((y_train, y_test))

train_size = 0.7
X_train, X_test, y_train, y_test = train_test_split(x, y, train_size=train_size)
X_train, X_test = torch.FloatTensor(X_train), torch.FloatTensor(X_test)
y_train, y_test = torch.LongTensor(y_train), torch.LongTensor(y_test)

# 展平图像
X_train = X_train.reshape(X_train.shape[0], -1)
X_test = X_test.reshape(X_test.shape[0], -1)

# 权重初始化：改用Xavier初始化，避免全0
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.xavier_normal_(m.weight)
        nn.init.zeros_(m.bias)

# 模型定义：移除最后一层ReLU，输出logits
model = nn.Sequential(
        nn.Linear(784, 200),
        nn.ReLU(),
        nn.Linear(200, 50),
        nn.ReLU(),
        nn.Linear(50, 10),  # 最后一层无激活，输出logits
    )

model = model.float()
model.apply(weights_init)

# 初始化母参数向量
mother_vector = nn.utils.parameters_to_vector(model.parameters())

# 超参数调整：适配MNIST的规模
SIGMA = 0.05
LR = 0.05
POPULATION_SIZE = 100
ITERATIONS = 200
BATCH_SIZE = 1024  # 用小批量计算fitness提升效率

# 损失与fitness函数：用负损失替代倒数，数值更稳定
loss_func = nn.CrossEntropyLoss()
def fitness_func(solution):
    nn.utils.vector_to_parameters(solution, model.parameters())
    # 随机采样小批量计算fitness，减少计算量
    idx = torch.randint(0, len(X_train), (BATCH_SIZE,))
    batch_x, batch_y = X_train[idx], y_train[idx]
    y_pred = model(batch_x)
    return -loss_func(y_pred, batch_y)  # 最大化负损失 = 最小化损失

# 生成带噪声的参数
def jitter(mother_params, noise):
    return mother_params + SIGMA * noise

# 计算种群fitness
def calculate_population_fitness(pop, mother_vector):
    fitness = torch.zeros(pop.shape[0])
    for i, noise in enumerate(pop):
        p_try = jitter(mother_vector, noise)
        fitness[i] = fitness_func(p_try)
    return fitness

# 测试准确率
def test(mother_params):
    nn.utils.vector_to_parameters(mother_params, model.parameters())
    with torch.no_grad():
        y_pred = torch.max(model(X_test), 1)[1]
        acc = (y_pred == y_test).sum().item() / len(y_test)
    return acc

n_params = mother_vector.shape[0]
print(f"Number of params: {n_params}")

# 训练循环
with torch.no_grad():
    for iteration in tqdm(range(ITERATIONS)):
        # 生成噪声种群
        pop = torch.randn(POPULATION_SIZE, n_params, dtype=torch.float32)
        fitness = calculate_population_fitness(pop, mother_vector)
        
        # 归一化fitness，处理标准差为0的特殊情况
        fitness_mean = torch.mean(fitness)
        fitness_std = torch.std(fitness)
        if fitness_std == 0:
            normalized_fitness = torch.zeros_like(fitness)
        else:
            normalized_fitness = (fitness - fitness_mean) / fitness_std
        
        # 更新母参数向量
        update = (LR / (POPULATION_SIZE * SIGMA)) * torch.matmul(pop.t(), normalized_fitness)
        mother_vector += update
        
        # 记录指标
        current_fitness = fitness_func(mother_vector)
        acc = test(mother_vector)
        print(f"Iteration: {iteration}, Fitness:{current_fitness:.4f}, Test Accuracy: {acc:.4f}")

额外优化建议 #

• GPU加速：将模型和数据移到GPU（model.to('cuda')、X_train = X_train.to('cuda')等），训练速度会大幅提升。
• 动态超参数调整：训练后期逐步减小SIGMA，让参数更新更精细。
• 并行计算：用PyTorch并行机制同时计算多个个体的fitness，进一步提升效率。

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