使用Deep Java Library训练图像文件夹数据集时遇形状不匹配错误

错误原因 #

错误提示的核心是模型输出维度与标签维度不匹配：

Exception in thread "main" ai.djl.engine.EngineException: MXNet engine call failed: MXNetError: Check failed: src.Size() == dst->Size() (20 vs. 200) : Cannot reshape array of size 20 into shape [20,10]

你当前使用的Loss.l2Loss()（均方误差损失）是为回归任务设计的，它要求目标标签和模型输出的维度完全一致（即每个样本对应10维的one-hot向量）。但DJL的ImageFolder数据集生成的标签是每个样本对应一个整数类别索引（形状为[batch_size]，这里是[20]），而模型输出是[batch_size, outputSize]（这里是[20,10]），两者维度不匹配导致报错。

解决方案 #

1. 替换为分类专用损失函数 #

将训练配置中的L2Loss替换为Softmax交叉熵损失，这是图像分类任务的标准损失函数，它可以直接处理整数类别索引标签：

修改代码中TrainingConfig的初始化部分：

// 修改后代码
TrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())                                      
        .addEvaluator(new Accuracy())                                                                  
        .optOptimizer(Optimizer.adadelta().build())                                                   
        .addTrainingListeners(TrainingListener.Defaults.logging());

2. 可选优化（推理阶段概率输出） #

如果后续需要模型输出类别概率，可以在模型最后一层添加Softmax激活：

SequentialBlock block = new SequentialBlock()                                                         
        .add(Blocks.batchFlattenBlock(inputSize))                                                     
        .add(Linear.builder().setUnits(inputSize).build())                                            
        .add(Activation::relu)                                                                         
        .add(Linear.builder().setUnits(128).build())                                                  
        .add(Activation::sigmoid)                                                                      
        .add(Linear.builder().setUnits(outputSize).build())
        .add(Activation::softmax); // 新增Softmax激活，用于输出概率

3. 验证数据集结构 #

确保你的数据集文件夹结构符合ImageFolder要求：

src/main/java/Ressources/
├── 类别1/
│   ├── img1.jpg
│   └── img2.jpg
├── 类别2/
│   ├── img3.jpg
│   └── ...
└── ...（共10个类别文件夹）

这样ImageFolder才能正确识别类别并生成对应标签。

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