keras实现注意力机制

一时失言乱红尘 2022-10-08 14:23 277阅读 0赞

class Attention_layer(Layer):
    
        def __init__(self,
                     W_regularizer=None, b_regularizer=None,
                     W_constraint=None, b_constraint=None,
                     bias=True, **kwargs):
    
            self.supports_masking = True
            self.init = initializers.get('glorot_uniform')
    
            self.W_regularizer = regularizers.get(W_regularizer)
            self.b_regularizer = regularizers.get(b_regularizer)
    
            self.W_constraint = constraints.get(W_constraint)
            self.b_constraint = constraints.get(b_constraint)
    
            self.bias = bias
            super(Attention_layer, self).__init__(**kwargs)
    
        def build(self, input_shape):
            assert len(input_shape) ==3
    
            self.W = self.add_weight((input_shape[-1], input_shape[-1],),
                                     initializer=self.init,
                                     name='{}_W'.format(self.name),
                                     regularizer=self.W_regularizer,
                                     constraint=self.W_constraint)
            if self.bias:
                self.b = self.add_weight((input_shape[-1],),
                                         initializer='zero',
                                         name='{}_b'.format(self.name),
                                         regularizer=self.b_regularizer,
                                         constraint=self.b_constraint)
    
            super(Attention_layer, self).build(input_shape)
    
        def compute_mask(self, input, input_mask=None):
            # do not pass the mask to the next layers
            return None
    
        def call(self, x, mask=None):
            uit = K.dot(x, self.W)
    
            if self.bias:
                uit += self.b
    
            uit = K.tanh(uit)
    
            a = K.exp(uit)
    
            # apply mask after the exp. will be re-normalized next
            if mask is not None:
                # Cast the mask to floatX to avoid float64 upcasting in theano
                a *= K.cast(mask, K.floatx())
    
            # in some cases especially in the early stages of training the sum may be almost zero
            # and this results in NaN's. A workaround is to add a very small positive number to the sum.
            # a /= K.cast(K.sum(a, axis=1, keepdims=True), K.floatx())
            a /= K.cast(K.sum(a, axis=1, keepdims=True) + K.epsilon(), K.floatx())
            print (a)
            # a = K.expand_dims(a)
            print (x)
            weighted_input = x * a
            print (weighted_input)
            return K.sum(weighted_input, axis=1)
    
        def compute_output_shape(self, input_shape):
            return (input_shape[0], input_shape[-1])

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9zaGVuY2h1bnh1LmJsb2cuY3Nkbi5uZXQ_size_1_color_FFFFFF_t_70][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9zaGVuY2h1bnh1LmJsb2cuY3Nkbi5uZXQ_size_1_color_FFFFFF_t_70]: /images/20221005/d0d854e497194522a00753528fc12bf5.png