解决Keras与TensorFlow版本不兼容的Attention Layer结构修改问题

Let's break down why your original code throws errors in TF2.x, then fix both versions of your Attention layer to work seamlessly with TensorFlow 2.2+ and Keras 2.4+.

Root Cause of the Error #

The main issues triggering the error are:

• TensorFlow 2.x enforces stricter weight naming checks. Using self.name directly to format weight names can lead to duplicate name conflicts when multiple instances of the layer are used.
• Custom layers in TF2.x must implement a proper get_config method to serialize all initialization parameters—this was missing in your first layer implementation and incomplete in the second.
• Minor API changes in Keras backend functions between TF1.x and TF2.x that break backward compatibility.

Fixed Version 1 (Bhadanau-style Attention with Step Dim Support) #

This fixes the naming conflict, adds proper serialization, and adapts to TF2.x conventions:

import tensorflow as tf
from tensorflow.keras.layers import Layer
from tensorflow.keras import initializers, regularizers, constraints

class Attention(Layer):
    '''Implement Attention is all you need's Bhadanau-style Attention'''
    def __init__(self, step_dim, 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
        self.step_dim = step_dim
        self.features_dim = 0
        super().__init__(**kwargs)  # Simplified super call for Python 3+

    def build(self, input_shape):
        assert len(input_shape) == 3
        # Let Keras handle weight naming automatically to avoid conflicts
        self.W = self.add_weight(shape=(input_shape[-1],),
                                 initializer=self.init,
                                 regularizer=self.W_regularizer,
                                 constraint=self.W_constraint,
                                 name='W')
        self.features_dim = input_shape[-1]
        if self.bias:
            self.b = self.add_weight(shape=(input_shape[1],),
                                     initializer='zero',
                                     regularizer=self.b_regularizer,
                                     constraint=self.b_constraint,
                                     name='b')
        else:
            self.b = None
        super().build(input_shape)  # Mark layer as built properly

    def compute_mask(self, inputs, mask=None):
        # Updated parameter name to match TF2.x conventions
        return None

    def call(self, x, mask=None):
        features_dim = self.features_dim
        step_dim = self.step_dim
        
        # Use TF2.x-compatible backend calls
        eij = tf.keras.backend.reshape(
            tf.keras.backend.dot(
                tf.keras.backend.reshape(x, (-1, features_dim)),
                tf.keras.backend.reshape(self.W, (features_dim, 1))
            ),
            (-1, step_dim)
        )
        if self.bias:
            eij += self.b
        eij = tf.keras.backend.tanh(eij)
        a = tf.keras.backend.exp(eij)
        
        if mask is not None:
            a *= tf.keras.backend.cast(mask, tf.keras.backend.floatx())
        
        a /= tf.keras.backend.cast(
            tf.keras.backend.sum(a, axis=1, keepdims=True) + tf.keras.backend.epsilon(),
            tf.keras.backend.floatx()
        )
        a = tf.keras.backend.expand_dims(a)
        weighted_input = x * a
        return tf.keras.backend.sum(weighted_input, axis=1)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.features_dim)
    
    def get_config(self):
        # Serialize all custom parameters for model saving/loading
        config = super().get_config()
        config.update({
            'step_dim': self.step_dim,
            'W_regularizer': regularizers.serialize(self.W_regularizer),
            'b_regularizer': regularizers.serialize(self.b_regularizer),
            'W_constraint': constraints.serialize(self.W_constraint),
            'b_constraint': constraints.serialize(self.b_constraint),
            'bias': self.bias
        })
        return config

Fixed Version 2 (Simplified Attention with MaxLen Support) #

This fixes the naming issue and completes the get_config implementation for your second layer:

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

class Attention(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def build(self, input_shape):
        # Let Keras handle weight naming to avoid conflicts
        self.W = self.add_weight(name="att_weight",
                                 shape=(input_shape[-1], 1),
                                 initializer="normal")
        self.b = self.add_weight(name="att_bias",
                                 shape=(input_shape[1], 1),
                                 initializer="zeros")
        super().build(input_shape)

    def call(self, x):
        et = tf.keras.backend.squeeze(
            tf.keras.backend.tanh(
                tf.keras.backend.dot(x, self.W) + self.b
            ),
            axis=-1
        )
        at = tf.keras.backend.softmax(et)
        at = tf.keras.backend.expand_dims(at, axis=-1)
        output = x * at
        return tf.keras.backend.sum(output, axis=1)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[-1])

    def get_config(self):
        # Pass super config (no custom params to serialize here)
        return super().get_config()

Key Modifications Explained #

1. Weight Naming: Removed manual self.name formatting in add_weight—Keras automatically adds the layer's name as a prefix to weights, eliminating duplicate name conflicts.
2. get_config Implementation: Added proper serialization of all custom parameters (for version 1) so you can save and load your model without issues in TF2.x.
3. TF2.x Compatibility: Used tf.keras.backend explicitly instead of the old standalone K, and updated method parameter names to match TF2.x conventions.
4. Simplified Super Calls: Used Python 3+ style super().__init__() for cleaner, more maintainable code.

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