TensorFlowJS多分类模型无法正常预测视线方向问题求助

1 核心架构错误修复 #

• 输出层激活函数选型错误：多分类任务输出层需要输出3个类别的概率分布，必须使用softmax激活函数，你当前用的tanh输出范围为[-1,1]，完全不符合分类交叉熵损失的输入要求，这是模型无法正常收敛的最核心原因。
• 输出层修改代码：

let congfig_output = {
    units: 3,
    activation: 'softmax', // 替换原来的tanh
}

• 损失函数配置补充：使用categoricalCrossentropy时可以新增准确率指标方便观察训练效果，同时调整学习率到常规初始值：

let config_compile = {
    optimizer: tf.train.adam(0.001), 
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy'] 
}

2 内存泄漏与浏览器崩溃问题修复 #

• 你当前代码中imageArray存储了大量未释放的张量，所有存在数组里的tf张量不会被tf.tidy自动回收，数据量变大后必然占用过多内存导致崩溃。
• 修复方案：
  • 收集数据时不要直接存储张量，先把张量转成普通数组存在内存里，训练前再统一转成张量，训练完成后立即释放大张量
  • tf.memory()需要打印才能看到输出，修改为console.log(tf.memory())即可看到当前张量占用、内存占用等数据
• 数据收集代码优化示例：

let collectData = async (label) => {
    const img = getImage();
    // 张量转普通数组存储，减少GPU显存占用
    const imgData = await img.array();
    imageArray.push(imgData);
    labelArray.push(label);
    img.dispose(); // 立即释放张量
}

3 训练逻辑优化（解决损失、准确率固定循环问题） #

• 数据集必须打乱：训练前调用tf.util.shuffleCombo同时打乱样本和标签，避免模型按样本顺序学习导致的伪收敛
• 超参数初始参考值：
  • 学习率：adam优化器初始用0.001即可，训练时如果损失震荡再下调到0.0005、0.0001
  • batch size：浏览器端训练建议设为8~16，不要超过32，避免显存不足
  • epoch：初始设20~30轮，观察验证集准确率停止上升就提前停止
• 样本均衡检查：确认三类样本的数量差不要超过2倍，避免模型偏向预测数量多的类别。

4 网络结构优化建议 #

• 单层卷积层特征提取能力不足，建议新增1~2组conv+pool层，提升模型对眼部特征的提取能力，参考结构：

let createModel = () => {
    const model = tf.sequential();
    // 第一组卷积池化
    model.add(tf.layers.conv2d({
        kernelSize: 3,
        filters: 16,
        strides: 1,
        activation: 'relu',
        inputShape: [imageHeight, imageWidth, imageChannels]
    }));
    model.add(tf.layers.maxPooling2d({poolSize: 2, strides: 2}));
    // 第二组卷积池化
    model.add(tf.layers.conv2d({
        kernelSize: 3,
        filters: 32,
        strides: 1,
        activation: 'relu'
    }));
    model.add(tf.layers.maxPooling2d({poolSize: 2, strides: 2}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dropout(0.3));
    // 全连接过渡层
    model.add(tf.layers.dense({units: 64, activation: 'relu'}));
    model.add(tf.layers.dropout(0.2));
    // 输出层
    model.add(tf.layers.dense({units: 3, activation: 'softmax'}));
    model.compile({
        optimizer: tf.train.adam(0.001),
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy']
    });
    return model;
}

5 其他细节优化 #

• 归一化逻辑可以调整为batchedImage.toFloat().div(tf.scalar(255))，输出范围[0,1]即可，不需要额外减1，适配relu激活的输入特性
• 1000张样本足够训练这个小模型，不需要再增加样本量，优先排查上面的逻辑错误即可。

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