未在自定义层的分配策略范围内创建变量答案

【问题标题】：Variable was not created in the distribution strategy scope with costum Layer未在自定义层的分配策略范围内创建变量
【发布时间】：2019-08-02 15:14:06
【问题描述】：

我正在尝试使用 TPU 微调 Google 协作实验室上的 BERT 模型。但我总是收到以下错误：

ValueError: 变量 (tf.Variable 'bert_layer_module/bert/encoder/layer_10/attention/output/LayerNorm/beta:0' shape=(768,) dtype=float32) 未在 (tensorflow) 的分布策略范围内创建.python.distribute.tpu_strategy.TPUStrategyV1 对象位于 0x7f6a1fad3390）。这很可能是由于并非所有层或模型或优化器都是在分发策略范围之外创建的。尝试确保您的代码类似于以下内容。

使用 strategy.scope()：

model=_create_model()

model.compile(...)

我的代码基于this notebook！我针对我的具体问题对其进行了修改，显然尝试在 TPU 上运行它。

我有一个服装层 BertLayer，它显然是在范围之外创建的：

class BertLayer(tf.keras.layers.Layer):
    def __init__(self, n_fine_tune_layers=10, **kwargs):
        self.n_fine_tune_layers = n_fine_tune_layers
        self.trainable = True
        self.output_size = 768
        super(BertLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.bert = hub.Module(
            bert_path,
            trainable=self.trainable,
            name="{}_module".format(self.name)
        )
        trainable_vars = self.bert.variables

        # Remove unused layers
        trainable_vars = [var for var in trainable_vars if not "/cls/" in var.name]

        # Select how many layers to fine tune
        trainable_vars = trainable_vars[-self.n_fine_tune_layers :]

        # Add to trainable weights
        for var in trainable_vars:
            self._trainable_weights.append(var)

        # Add non-trainable weights
        for var in self.bert.variables:
            if var not in self._trainable_weights:
                self._non_trainable_weights.append(var)

        super(BertLayer, self).build(input_shape)

    def call(self, inputs):
        inputs = [K.cast(x, dtype="int32") for x in inputs]
        input_ids, input_mask, segment_ids = inputs
        bert_inputs = dict(
            input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids
        )
        result = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
            "pooled_output"
        ] 
        return result

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

模型创建在这里完成：

def build_model(max_seq_length): 
    output_classes = train_labels[0].shape
    # Build model
    in_id = tf.keras.layers.Input(shape=(max_seq_length,), name="input_ids")
    in_mask = tf.keras.layers.Input(shape=(max_seq_length,), name="input_masks")
    in_segment = tf.keras.layers.Input(shape=(max_seq_length,), name="segment_ids")
    bert_inputs = [in_id, in_mask, in_segment]

    # Instantiate the custom Bert Layer defined above
    bert_output = BertLayer(n_fine_tune_layers=10)(bert_inputs)

    # Build the rest of the classifier 
    dense = tf.keras.layers.Dense(256, activation='relu')(bert_output)
    pred = tf.keras.layers.Dense(train_labels.shape[1], activation='sigmoid')(dense)

    model = tf.keras.models.Model(inputs=bert_inputs, outputs=pred)
    return model

调用model.compile时出错

strategy = tf.distribute.experimental.TPUStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS))
with strategy.scope():
  model = build_model(256)

  opt = tf.train.AdamOptimizer(0.001)
  model.compile(loss='binary_crossentropy', optimizer=opt)

据我了解，BertLayer 确实是在范围内创建的，但我对 keras 和 tensorflow 比较陌生，所以很高兴得到您的帮助。我正在开发 tensorflow 1.14

【问题讨论】：

标签： python tensorflow keras google-colaboratory tpu

【解决方案1】：

考虑在策略范围内指定模型输入的形状，您可以尝试以下方法之一： 1.模型创建后拨打model.build()。 2.在__init__()中定义模型第一层的input_shape。 3. 创建模型后，使用实数张量（ones, zeros, ect.）调用model(tensor)。

无论如何，指定模型输入的形状。这对我有用。希望这会对你有所帮助。

【讨论】：

【解决方案2】：

对我有用的是在 with tpu_strategy.scope(): 块中完成整个模型定义（不仅仅是 model = build...() 语句）。

我使用的是 Keras 功能 API，我想我可以像往常一样定义模型层，然后只将 model = keras.Model(inputs=model_inputs, outputs=model_outputs) 放在 with tpu_strategy.scope(): 块中。然而这并没有奏效。

在您的具体情况下，这会将class BertLayer(tf.keras.layers.Layer): 和最终模型中包含的所有其他函数/类/等放在with tpu_strategy.scope(): 块中。

【讨论】：