如何重新配置YOLO2输出层输出特定维度？DL4J自定义层问询

API 探路者

2026-5-14

如何重新配置YOLO2输出层以适配特定维度？

问题背景

各位DL4J开发者，能否实现一种非默认5参数[x, y, w, h, p]的YOLO变体？以下是我从dl4j-examples官方仓库获取的、针对13×13网格锁定图像的代码：

...
graphBuilder.addLayer("convolution2d_23", new ConvolutionLayer.Builder(1,1)
.nIn(1024)
.nOut(nBoxes* (5 +nClasses))//don't want the 5 default params always
.weightInit(WeightInit.XAVIER)
.stride(1,1)
.convolutionMode(ConvolutionMode.Same)
.weightInit(WeightInit.RELU)
.activation(Activation.IDENTITY)
.cudnnAlgoMode(cudnnAlgoMode)
.build(), "activation_22")
.addLayer("outputs", new Yolo2OutputLayer.Builder()
.boundingBoxPriors(priors)
.build(), "convolution2d_23")
.setOutputs("outputs");
graphBuilder.build();
...

我需要对Yolo2OutputLayer进行重新配置，或自定义Yolo2OutputLayer类，使其能够输出任意特定维度。目前我需要输出13×13×80的维度，其中每个1×1×80的单元切片对应1×1×2[x, y, w, h, p, c, a0, a1,..., a31]，其中2为每个网格单元的边界框数量：

x - 边界框x坐标 - 1维
y - 边界框y坐标 - 1维
w - 边界框宽度 - 1维
h - 边界框高度 - 1维
p - 边界框目标置信度 - 1维
c - 边界框类别（3类） - 3维
a0...a31 - 自定义参数 - 32维的内容

解决方案

要实现这个自定义输出的YOLO2变体，核心是自定义Yolo2OutputLayer类——毕竟DL4J默认的Yolo2OutputLayer是硬编码针对5个框参数+类别的逻辑。下面一步步来实现：

1. 先调整卷积层的输出通道数

按照你的需求，每个边界框的总参数维度是1(x)+1(y)+1(w)+1(h)+1(p)+3(c)+32(a) = 40维，2个框的话总通道数就是2*40=80，正好对应13×13×80的输出。所以卷积层的nOut要改成：

nOut(nBoxes * (5 + 3 + 32)) // 5是基础框参数，3是类别数，32是自定义参数

2. 自定义Yolo2OutputLayer子类

默认的Yolo2OutputLayer只处理框参数和类别概率，我们需要扩展它来支持自定义参数，主要修改损失计算和输出激活两个核心逻辑：

public class CustomYolo2OutputLayer extends Yolo2OutputLayer {

    // 自定义参数的维度数，这里设为32
    private int customParamsSize = 32;
    // 类别数，这里设为3
    private int numClasses = 3;

    @Override
    public void computeGradientAndScore(INDArray input, INDArray labels, INDArray gradient, INDArray score) {
        int numBoxes = getBoundingBoxPriors().length / 2;
        int paramsPerBox = 5 + numClasses + customParamsSize;

        // 1. 拆分输入张量的各个部分：基础框参数、类别、自定义参数
        int baseEndIdx = numBoxes * (5 + numClasses);
        INDArray parentInput = input.get(NDArrayIndex.all(), NDArrayIndex.interval(0, baseEndIdx), NDArrayIndex.all(), NDArrayIndex.all());
        INDArray customParams = input.get(NDArrayIndex.all(), NDArrayIndex.interval(baseEndIdx, numBoxes*paramsPerBox), NDArrayIndex.all(), NDArrayIndex.all());

        // 2. 复用父类逻辑处理框参数和类别的损失
        INDArray parentLabels = labels.get(NDArrayIndex.all(), NDArrayIndex.interval(0, baseEndIdx), NDArrayIndex.all(), NDArrayIndex.all());
        super.computeGradientAndScore(parentInput, parentLabels, gradient.get(NDArrayIndex.all(), NDArrayIndex.interval(0, baseEndIdx), NDArrayIndex.all(), NDArrayIndex.all()), score);

        // 3. 添加自定义参数的损失计算（这里用MSE损失示例，可按需替换）
        INDArray customLabels = labels.get(NDArrayIndex.all(), NDArrayIndex.interval(baseEndIdx, numBoxes*paramsPerBox), NDArrayIndex.all(), NDArrayIndex.all());
        INDArray customLoss = LossMSE.computeLoss(customParams, customLabels, null, true);
        score.addi(customLoss);

        // 计算自定义参数的梯度
        INDArray customGradient = gradient.get(NDArrayIndex.all(), NDArrayIndex.interval(baseEndIdx, numBoxes*paramsPerBox), NDArrayIndex.all(), NDArrayIndex.all());
        LossMSE.computeGradient(customParams, customLabels, customGradient);
    }

    @Override
    public INDArray activate(boolean training) {
        INDArray output = super.activate(training);
        int numBoxes = getBoundingBoxPriors().length / 2;
        int customStartIdx = numBoxes*(5 + numClasses);
        
        // 如果自定义参数需要特定激活（比如Sigmoid归一化到0-1），在这里处理
        // 示例：保持线性激活，无需处理则直接返回output
        INDArray customPart = output.get(NDArrayIndex.all(), NDArrayIndex.interval(customStartIdx, output.size(1)), NDArrayIndex.all(), NDArrayIndex.all());
        return output;
    }

    // 提供setter方便配置参数
    public void setCustomParamsSize(int customParamsSize) {
        this.customParamsSize = customParamsSize;
    }

    public void setNumClasses(int numClasses) {
        this.numClasses = numClasses;
    }
}

3. 替换原代码中的输出层

在模型构建代码里，把默认的Yolo2OutputLayer换成自定义类，并配置对应参数：

graphBuilder.addLayer("convolution2d_23", new ConvolutionLayer.Builder(1,1)
        .nIn(1024)
        .nOut(2 * (5 + 3 + 32)) // 2个框，每个框40维参数
        .weightInit(WeightInit.XAVIER)
        .stride(1,1)
        .convolutionMode(ConvolutionMode.Same)
        .activation(Activation.IDENTITY)
        .cudnnAlgoMode(cudnnAlgoMode)
        .build(), "activation_22")
.addLayer("outputs", new CustomYolo2OutputLayer.Builder()
        .boundingBoxPriors(priors)
        .numClasses(3)
        .customParamsSize(32)
        .build(), "convolution2d_23")
.setOutputs("outputs");
graphBuilder.build();

4. 调整训练标签格式

训练时的标签张量要和输出结构对齐：每个网格单元的每个框，标签需要包含x,y,w,h,p,3维类别标签,32维自定义参数，总维度保持13×13×80。

关键注意事项

损失平衡：自定义参数的损失量级可能和框/类别的损失差异很大，建议给不同部分损失添加权重系数，避免某部分损失主导训练。
激活选择：如果自定义参数有范围限制（比如0-1），记得在activate方法里添加对应激活处理；无范围限制的回归值用线性激活即可。
参数初始化：自定义参数对应的卷积层权重，建议用Xavier或He初始化，避免训练初期梯度爆炸/消失。

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