自定义TensorFlow模型权重加载失败求助（附Seq2Seq代码）

环境信息 #

• TensorFlow版本：2.16.1
• Keras版本：3.3.3
• 运行环境：Google Colab

报错信息 #

ValueError                                Traceback (most recent call last)
<ipython-input-46-4d65c8f7b80b> in <cell line: 1>()
----> 1 m3.load_weights("./model_v8.weights.h5")

1 frames
/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py in error_handler(*args, **kwargs)
    120             # To get the full stack trace, call:
    121             # `keras.config.disable_traceback_filtering()`
---> 122             raise e.with_traceback(filtered_tb) from None
    123         finally:
    124             del filtered_tb

/usr/local/lib/python3.10/dist-packages/keras/src/saving/saving_lib.py in _raise_loading_failure(error_msgs, warn_only)
    293         warnings.warn(msg)
    294     else:
---> 295         raise ValueError(msg)
    296 
    297 

ValueError: A total of 1 objects could not be loaded. Example error message for object <Embedding name=embedding_10, built=True>:

Layer 'embedding_10' expected 1 variables, but received 0 variables during loading. Expected: ['embeddings']

List of objects that could not be loaded:
[<Embedding name=embedding_10, built=True>]

模型代码 #

Encoder #

class Encoder(Layer):
    def __init__(self,
                tokenizer,
                embedding_size,
                hidden_units):
        """
            Encoder Block in seq2seq

        :param tokenizer: tokenizer of the source language
        :param embedding_size: dimensionality of the embedding layer
        :param hidden_units: dimensionality of the output
        """

        super().__init__()
        self.tokenizer = tokenizer
        self.embedding_size = embedding_size
        self.hidden_units = hidden_units
        self.vocab_size = tokenizer.vocabulary_size()
        self.embedding = Embedding(input_dim=self.vocab_size,
                                   output_dim=embedding_size)
        self.rnn = Bidirectional(
            merge_mode="sum",
            layer=LSTM(units=hidden_units,
                    dropout=DROPOUT,
                    return_sequences=True,
                    return_state=True))

    def call(self,
            x,
            training=True):
        """
        :param x: [batch, time_steps]
        :param training: is training or not
        :return:
            encoder_hidden_state: [batch, hidden_state_dim]
            state_h: [batch, hidden_state_dim]
            state_c: [batch, hidden_state_dim]
        """
        mask = tf.where(x != 0, True, False)
        x = self.embedding(x)
        x, forward_h, forward_c, backward_h, backward_c = self.rnn(x, mask=mask,
                                                                training=training)

        return x, forward_h + backward_h, forward_c + backward_c

Bahdanau Attention #

class BahdanauAttention(Layer):
    def __init__(self,
                 hidden_units):
        super().__init__()
        self.Va = Dense(1)
        self.Wa = Dense(hidden_units)
        self.Ua = Dense(hidden_units)
        self.norm = LayerNormalization()
        self.tanh = Activation(tf.keras.activations.tanh)
        self.add = Add()

    def call(self,
             context, x):
        """
            Calculate the context vector based on all encoder hidden states and
            previous decoder state.

        :param: context: tensor, all encoder hidden states
        :param: x: tensor, previous state from Decoder
        :return:
            context_vector: tensor, the calculated context vector based on the
            input parameters
        """
        # Expand dims to ensure scores shape = [batch, Ty, Tx]
        context = tf.expand_dims(context, axis=1)
        x = tf.expand_dims(x, axis=2)

        scores = self.Va(self.tanh(self.add([self.Wa(context), self.Ua(x)])))
        scores = tf.squeeze(scores, axis=-1)
        attn_weights = tf.nn.softmax(scores, axis=-1)

        # NOTE: context shape = [batch, 1, Tx, feature] so that expand
        # dim of attention weights
        context_vector = tf.expand_dims(attn_weights, axis=-1) * context
        context_vector = tf.reduce_sum(context_vector, axis=-2)
        context_vector = self.norm(context_vector)
        context_vector = self.add([context_vector, tf.squeeze(x, -2)])

        return context_vector

Decoder #

class Decoder(Layer):
    def __init__(self,
                tokenizer,
                embedding_size,
                hidden_units):
        """
            Decoder Block in seq2seq

        :param tokenizer: tokenizer of the source language
        :param embedding_size: dimensionality of the embedding layer
        :param hidden_units: dimensionality of the output
        """

        super().__init__()
        self.tokenizer = tokenizer
        self.embedding_size = embedding_size
        self.hidden_units = hidden_units
        self.vocab = tokenizer.get_vocabulary()
        self.vocab_size = tokenizer.vocabulary_size()
        self.embedding = Embedding(input_dim=self.vocab_size,
                                output_dim=embedding_size)
        self.rnn = LSTM(units=hidden_units,
                        dropout=DROPOUT,
                        return_sequences=True,
                        return_state=True)
        self.attention = BahdanauAttention(hidden_units)
        self.dense = Dense(self.vocab_size)

    def call(self,
            context, x,
            encoder_state,
            training=True,
            return_state=False):
        """
        :param context: all encoder states
        :param x: all initial decoder states
        :param encoder_state: last state from encoder
        :param training:
        :param return_state:

        :return:
            logits:
            state_h: hidden state
            state_c: cell state
        """
        mask = tf.where(x != 0, True, False)
        x = self.embedding(x)
        decoder_outputs, state_h, state_c = self.rnn(x, initial_state=encoder_state,
                                                    mask=mask,
                                                    training=training)
        dense_inputs = self.attention(context, decoder_outputs)
        logits = self.dense(dense_inputs)

        if return_state:
            return logits, state_h, state_c
        else:
            return logits

NMT Model #

class NMT(Model):
    @classmethod
    def add_method(cls, fun):
        setattr(cls, fun.__name__, fun)
        return fun

    def __init__(self,
                 input_tokenizer,
                 output_tokenizer,
                 embedding_size,
                 hidden_units):
        """
            Initialize an instance for Neural Machine Translation Task

        :param input_tokenizer: tokenizer of the input language
        :param output_tokenizer: tokenizer of the output language
        :param embedding_size: dimensionality of embedding layer
        :param hidden_units: dimensionality of the output
        """

        super().__init__()
        self.input_tokenizer = input_tokenizer
        self.output_tokenizer = output_tokenizer
        self.embedding_size = embedding_size
        self.hidden_units = hidden_units
        self.encoder = Encoder(input_tokenizer,
                               embedding_size,
                               hidden_units)
        self.decoder = Decoder(output_tokenizer,
                               embedding_size,
                               hidden_units)

    def call(self,
             inputs):
        encoder_inputs, decoder_inputs = inputs
        encoder_outputs, state_h, state_c = self.encoder(encoder_inputs)
        logits = self.decoder(encoder_outputs, decoder_inputs,
                              [state_h, state_c])

        return logits

    @NMT.add_method
    def translate(self, next_inputs,
                maxlen=40):
        """
        """
        def sampling(logits):
            probs = tf.nn.softmax(logits)
            dist = probs.numpy().squeeze()
            idx = np.random.choice(range(self.decoder.vocab_size), p=dist)

            return idx

        translation = []
        next_inputs = expand_contractions(next_inputs.lower(), en_contraction_map)
        next_idx = np.asarray(self.encoder.tokenizer(next_inputs))

        while next_idx.ndim != 2:
            next_idx = tf.expand_dims(next_idx, axis=0)

        encoder_outputs, state_h, state_c = self.encoder(next_idx, training=False)

        next_inputs = "[START]"
        next_idx = np.asarray(word_to_idx[next_inputs])

        for i in range(maxlen):
            while next_idx.ndim != 2:
                next_idx = tf.expand_dims(next_idx, axis=0)

            logits, state_h, state_c = self.decoder(encoder_outputs, next_idx,
                                                    [state_h, state_c],
                                                    training=False,
                                                    return_state=True)
            next_idx = sampling(logits)
            next_inputs = self.decoder.vocab[next_idx]

            if next_inputs == "[END]":
                break
            elif next_inputs == "[UNK]":
                continue
            else:
                translation.append(next_inputs)

        return " ".join(translation)

已尝试的操作 #

• 确认TensorFlow/Keras版本与权重文件生成环境一致
• 尝试保存整个模型并添加get_config()/from_config()方法，仍报错且收到Decoder类build方法警告，转回仅加载权重的方式
• 尝试先调用translate()方法初始化层，再加载权重，依然失败

解决方案 #

一、解决权重加载失败问题 #

1. 确保模型结构与权重完全匹配

• 加载权重时的模型实例，所有参数（vocab_size、embedding_size、hidden_units）必须与保存权重时的模型完全一致，且tokenizer的词汇表需完全相同（避免Embedding层输入维度变化）。
• 检查word_to_idx映射是否与output_tokenizer的词汇表一致，translate方法中使用的word_to_idx["[START]"]若与训练时的映射不匹配，会导致Decoder的Embedding层初始化异常。

2. 完整初始化模型变量后再加载权重

实例化模型后，先传入符合输入形状的虚拟数据执行一次前向传播，确保所有层（包括Embedding、LSTM等）的变量被完整初始化，再加载权重：

# 生成虚拟输入，形状与实际训练数据一致
dummy_encoder_input = tf.random.uniform((1, 20), 0, input_tokenizer.vocabulary_size(), dtype=tf.int32)
dummy_decoder_input = tf.random.uniform((1, 20), 0, output_tokenizer.vocabulary_size(), dtype=tf.int32)

# 执行前向传播，初始化所有层变量
m3([dummy_encoder_input, dummy_decoder_input])

# 加载权重
m3.load_weights("./model_v8.weights.h5")

这种方式比调用translate()更可靠，能确保所有层都完成初始化。

3. 检查权重保存方式

确保保存权重时使用的是model.save_weights("xxx.weights.h5")，而非其他保存格式（如完整模型的H5文件结构不同）。

二、自定义Layer的build方法实现 #

对于由内置Keras Layer组成的自定义Layer，通常无需手动实现build方法——内置Layer会在第一次call时自动完成构建。若需要手动控制构建逻辑，或包含自定义变量，可按以下方式实现：

示例：为Decoder类添加build方法

class Decoder(Layer):
    def __init__(self, tokenizer, embedding_size, hidden_units):
        super().__init__()
        self.tokenizer = tokenizer
        self.embedding_size = embedding_size
        self.hidden_units = hidden_units
        self.vocab = tokenizer.get_vocabulary()
        self.vocab_size = tokenizer.vocabulary_size()
        # 仅定义层的占位，不在__init__中实例化
        self.embedding = None
        self.rnn = None
        self.attention = None
        self.dense = None

    def build(self, input_shape):
        # 根据输入形状初始化所有子层
        self.embedding = Embedding(input_dim=self.vocab_size,
                                   output_dim=self.embedding_size)
        self.rnn = LSTM(units=self.hidden_units,
                        dropout=DROPOUT,
                        return_sequences=True,
                        return_state=True)
        self.attention = BahdanauAttention(self.hidden_units)
        self.dense = Dense(self.vocab_size)
        super().build(input_shape)  # 必须调用父类build方法完成注册

    def call(self, context, x, encoder_state, training=True, return_state=False):
        # 确保层已构建
        if not self.built:
            self.build([context.shape, x.shape])
        mask = tf.where(x != 0, True, False)
        x = self.embedding(x)
        decoder_outputs, state_h, state_c = self.rnn(x, initial_state=encoder_state,
                                                    mask=mask,
                                                    training=training)
        dense_inputs = self.attention(context, decoder_outputs)
        logits = self.dense(dense_inputs)

        if return_state:
            return logits, state_h, state_c
        else:
            return logits

若需支持模型完整序列化（保存/加载整个模型）

需为每个自定义Layer实现get_config和from_config方法，注意Keras Tokenizer无法直接序列化，需单独保存词汇表后重建：

class Encoder(Layer):
    # ... 原有代码 ...

    def get_config(self):
        config = super().get_config()
        config.update({
            "embedding_size": self.embedding_size,
            "hidden_units": self.hidden_units,
            "vocab_size": self.vocab_size
            # 不直接保存tokenizer实例，而是保存词汇表信息
        })
        return config

    @classmethod
    def from_config(cls, config):
        # 加载时需重新创建tokenizer，传入对应的词汇表
        vocab_size = config.pop("vocab_size")
        # 假设已从文件加载词汇表并创建tokenizer
        tokenizer = create_tokenizer_from_vocab(vocab_path)
        return cls(tokenizer, **config)

内容的提问来源于stack exchange，提问作者Kỳ Nguyễn Cao