基于Jax Flax复现PyTorch风格MLP时遇形状不兼容错误求助

核心错误原因 #

1. 未扁平化图像输入：MNIST输入是(批量大小, 28, 28, 1)的4维张量，Flax的nn.Dense仅作用于最后一维，若不先扁平化为(批量大小, 784)，会导致每一层输出仍保留空间维度（如(128,28,28,50)），最终logits维度与one-hot标签（(128,10)）不匹配，引发广播错误。
2. 模型结构重复冗余：原代码循环中创建了hidden_size长度的Dense层，之后又额外添加了Dense(self.hidden_size[-1])，导致输出维度错误；同时未对齐PyTorch的层结构（PyTorch是len(layer)-1层Linear，对应输入到输出的映射）。
3. Training参数设计不合理：将training作为类属性传入，不符合Flax的范式——Flax通常通过deterministic参数或rngs控制Dropout的训练/推理模式，而非类实例化时固定状态。
4. Train_step中的不当操作：在loss_fn内重复实例化模型，且固定Dropout的随机种子，既浪费计算资源，也无法保证训练过程的随机性。

修正后的代码 #

1. 正确的MLP模型定义（对齐PyTorch逻辑） #

from flax import linen as nn
from typing import Sequence

class MLPModel(nn.Module):
    hidden_sizes: Sequence[int]
    dp_rate: float

    @nn.compact
    def __call__(self, x, deterministic=False):
        # 先扁平化输入：(batch, 28, 28, 1) -> (batch, 784)
        x = x.reshape(x.shape[0], -1)
        for idx in range(len(self.hidden_sizes) - 1):
            x = nn.Dense(self.hidden_sizes[idx+1])(x)
            x = nn.relu(x)
            x = nn.Dropout(self.dp_rate)(x, deterministic=deterministic)
        # 最后一层不加Dropout，直接输出log_softmax
        x = nn.Dense(self.hidden_sizes[-1])(x)
        x = nn.log_softmax(x, axis=-1)
        return x

2. 修正后的Train_step函数 #

import jax
import jax.numpy as jnp
from flax.training import train_state
import optax

@jax.jit
def train_step(state, imgs, gt_labels, key):
    def loss_fn(params):
        # 使用传入的key生成Dropout子密钥，避免固定种子
        dropout_key = jax.random.fold_in(key, state.step)
        logits = state.apply_fn(
            params, 
            imgs, 
            deterministic=False,
            rngs={'dropout': dropout_key}
        )
        one_hot_gt_labels = jax.nn.one_hot(gt_labels, num_classes=10)
        # 用optax内置交叉熵计算更简洁稳定
        loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot_gt_labels).mean()
        return loss, logits
    
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    metrics = compute_metrics(logits=logits, gt_labels=gt_labels)
    return state, metrics

3. 模型初始化示例 #

# 初始化模型与训练状态
model = MLPModel(hidden_sizes=[784, 50, 50, 10], dp_rate=0.1)
key = jax.random.PRNGKey(42)
params = model.init(key, jnp.ones((1, 28, 28, 1)))['params']
tx = optax.adam(learning_rate=1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

关键说明 #

• 输入扁平化：必须将MNIST的28x28图像展平为784维向量，才能让Dense层正确处理特征。
• Dropout控制：通过deterministic参数区分训练（False）和推理（True）模式，同时使用fold_in生成随step变化的随机密钥，保证训练随机性。
• 层结构对齐：循环次数设为len(hidden_sizes)-1，对应从输入维度到中间层再到输出层的完整映射，避免重复添加Dense层。

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