keras

Keras中如何定义一个自定义的层

小樊
114
2024-03-14 10:54:29
栏目: 深度学习

要定义一个自定义的层,需要继承keras.layers.Layer类,并重写__init__call方法。下面是一个简单的示例:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Layer

class MyCustomLayer(Layer):
    def __init__(self, output_dim, activation=None, **kwargs):
        self.output_dim = output_dim
        self.activation = keras.activations.get(activation)
        super(MyCustomLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel', 
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)
        super(MyCustomLayer, self).build(input_shape)

    def call(self, inputs):
        output = tf.matmul(inputs, self.kernel)
        if self.activation is not None:
            output = self.activation(output)
        return output

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim)

在这个示例中,我们定义了一个自定义的层MyCustomLayer,它具有一个可调节的输出维度和激活函数。在__init__方法中设置了输出维度和激活函数,并在build方法中创建了权重矩阵。在call方法中实现了层的前向传播逻辑,并在最后返回输出。最后,compute_output_shape方法用于计算输出的形状。

定义好自定义的层后,可以像使用其他内置的层一样将其添加到模型中进行训练。

0
看了该问题的人还看了