Run this notebook online:Binder or Colab: Colab

5.3. 自定义层

深度学习成功背后的一个因素是,可以用创造性的方式组合广泛的层,从而设计出适用于各种任务的结构。例如,研究人员发明了专门用于处理图像、文本、序列数据和执行动态编程的层。早晚有一天,你会遇到或要自己发明一个在深度学习框架中还不存在的层。在这些情况下,你必须构建自定义层。在本节中,我们将向你展示如何操作。

5.3.1. 不带参数的层

首先,我们构造一个没有任何参数的自定义层。如果你还记得我们在 Section 5.1 对块的介绍,这应该看起来很眼熟。下面的 CenteredLayer 类要从其输入中减去均值。要构建它,我们只需继承基础层类并实现正向传播功能。

%load ../utils/djl-imports
class CenteredLayer extends AbstractBlock {

    public CenteredLayer() {
        super((byte)1);
    }

    @Override
    protected NDList forwardInternal(
            ParameterStore parameterStore,
            NDList inputs,
            boolean training,
            PairList<String, Object> params) {
        NDList current = inputs;
        // Subtract the mean from the input
        return new NDList(current.head().sub(current.head().mean()));
    }

    @Override
    public Shape[] getOutputShapes(Shape[] inputs) {
        // Output shape should be the same as input
        return inputs;
    }
}

让我们通过向其提供一些数据来验证该层是否按预期工作。

NDManager manager = NDManager.newBaseManager();

CenteredLayer layer = new CenteredLayer();

Model model = Model.newInstance("centered-layer");
model.setBlock(layer);

Predictor<NDList, NDList> predictor = model.newPredictor(new NoopTranslator());
NDArray input = manager.create(new float[]{1f, 2f, 3f, 4f, 5f});
predictor.predict(new NDList(input)).singletonOrThrow();
ND: (5) gpu(0) float32
[-2., -1.,  0.,  1.,  2.]

现在,我们可以将层作为组件合并到构建更复杂的模型中。

SequentialBlock net = new SequentialBlock();
net.add(Linear.builder().setUnits(128).build());
net.add(new CenteredLayer());
net.setInitializer(new NormalInitializer(), Parameter.Type.WEIGHT);
net.initialize(manager, DataType.FLOAT32, input.getShape());

作为额外的健全性检查,我们可以向网络发送随机数据后,检查均值是否为0。由于我们处理的是浮点数,因为存储精度的原因,我们仍然可能会看到一个非常小的非零数。

NDArray input = manager.randomUniform(-0.07f, 0.07f, new Shape(4, 8));
NDArray y = predictor.predict(new NDList(input)).singletonOrThrow();
y.mean();
ND: () gpu(0) float32
 6.98491931e-10

5.3.2. 带参数的图层

既然我们知道了如何定义简单的层,那么让我们继续定义具有参数的层,这些参数可以通过训练进行调整。我们可以使用内置函数来创建参数,这些参数提供一些基本的管理功能。比如管理访问、初始化、共享、保存和加载模型参数。这样做的好处之一是,我们不需要为每个自定义层编写自定义序列化程序。

现在,让我们实现自定义版本的全连接层。回想一下,该层需要两个参数,一个用于表示权重,另一个用于表示偏置项。在此实现中,我们使用 Activation.relu 作为激活函数。该层需要输入参数:inUnitsoutUnits,分别表示输入和输出的数量。

class MyLinear extends AbstractBlock {

    private Parameter weight;
    private Parameter bias;

    private int inUnits;
    private int outUnits;

    // outUnits: the number of outputs in this layer
    // inUnits: the number of inputs in this layer
    public MyLinear(int outUnits, int inUnits) {
        super((byte)1);
        this.inUnits = inUnits;
        this.outUnits = outUnits;
        weight = addParameter(
            Parameter.builder()
                .setName("weight")
                .setType(Parameter.Type.WEIGHT)
                .optShape(new Shape(inUnits, outUnits))
                .build());
        bias = addParameter(
            Parameter.builder()
                .setName("bias")
                .setType(Parameter.Type.BIAS)
                .optShape(new Shape(outUnits))
                .build());
    }

    @Override
    protected NDList forwardInternal(
            ParameterStore parameterStore,
            NDList inputs,
            boolean training,
            PairList<String, Object> params) {
        NDArray input = inputs.singletonOrThrow();
        Device device = input.getDevice();
        // Since we added the parameter, we can now access it from the parameter store
        NDArray weightArr = parameterStore.getValue(weight, device, false);
        NDArray biasArr = parameterStore.getValue(bias, device, false);
        return relu(linear(input, weightArr, biasArr));
    }

    @Override
    public Shape[] getOutputShapes(Shape[] inputs) {
        return new Shape[]{new Shape(outUnits, inUnits)};
    }

    // Applies linear transformation
    private static NDArray linear(NDArray input, NDArray weight, NDArray bias) {
        return input.dot(weight).add(bias);
    }

    // Applies relu transformation
    private static NDList relu(NDArray input) {
        return new NDList(Activation.relu(input));
    }
}

接下来,我们实例化MyLinear类并访问其模型参数。

// 5 units in -> 3 units out
MyLinear linear = new MyLinear(3, 5);
var params = linear.getParameters();
for (Pair<String, Parameter> param : params) {
    System.out.println(param.getKey());
}
weight
bias

我们可以使用自定义层直接执行正向传播计算。

NDArray input = manager.randomUniform(0, 1, new Shape(2, 5));

linear.initialize(manager, DataType.FLOAT32, input.getShape());

Model model = Model.newInstance("my-linear");
model.setBlock(linear);

Predictor<NDList, NDList> predictor = model.newPredictor(new NoopTranslator());
predictor.predict(new NDList(input)).singletonOrThrow();
ND: (2, 3) gpu(0) float32
[[ 6.66518092e-01,  0.00000000e+00,  1.33861804e+00],
 [ 0.00000000e+00,  1.01035833e-03,  2.03937054e-01],
]

我们还可以使用自定义层构建模型。我们可以像使用内置的全连接层一样使用自定义层。

NDArray input = manager.randomUniform(0, 1, new Shape(2, 64));

SequentialBlock net = new SequentialBlock();
net.add(new MyLinear(8, 64)); // 64 units in -> 8 units out
net.add(new MyLinear(1, 8)); // 8 units in -> 1 unit out
net.initialize(manager, DataType.FLOAT32, input.getShape());

Model model = Model.newInstance("lin-reg-custom");
model.setBlock(net);

Predictor<NDList, NDList> predictor = model.newPredictor(new NoopTranslator());
predictor.predict(new NDList(input)).singletonOrThrow();
ND: (2, 1) gpu(0) float32
[[0.],
 [0.],
]

5.3.3. 小结

  • 我们可以通过基本层类设计自定义层。这允许我们定义灵活的新层,其行为与库中的任何现有层不同。

  • 在自定义层定义完成后,就可以在任意环境和网络结构中调用该自定义层。

  • 层可以有局部参数,这些参数可以通过内置函数创建。

5.3.4. 练习

  1. 设计一个接受输入并计算 NDArray 汇总的层,它返回\(y_k = \sum_{i, j} W_{ijk} x_i x_j\)

  2. 设计一个返回输入数据的傅立叶系数前半部分的层。