Tensorflow 2.1 TPU v2 使用 bfloat16 减少内存使用

Question

我对 TPUv2 的内存使用有一些疑问。 我想对一些大型模型进行一些实验，但不幸的是该模型不适合内存。我想使用 bfloat16 来节省一些内存，但是在调用模型时遇到了一些问题：

try:
    # TPU detection. No parameters necessary if TPU_NAME environment variable is
    # set: this is always the case on Kaggle.
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    print('Running on TPU ', resolver.master())
except ValueError:
    resolver = None

if resolver:

  tf.config.experimental_connect_to_cluster(resolver)
  tf.tpu.experimental.initialize_tpu_system(resolver)
  strategy = tf.distribute.experimental.TPUStrategy(resolver)
else:
    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.
    strategy = tf.distribute.get_strategy()


policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16')
tf.keras.mixed_precision.experimental.set_policy(policy)

with strategy.scope():

    model = CustomModel(TFXLMRobertaModel.from_pretrained("jplu/tf-xlm-roberta-large"),  num_classes=5)
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)
    optimizer = tf.mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')

    model.compile(optimizer=optimizer,loss=['mse']) 

InvalidArgumentError Traceback（最近调用 最后）

在 () 3 与 strategy.scope()： 4 ----> 5 模型 = CustomModel(TFXLMRobertaModel.from_pretrained("jplu/tf-xlm-roberta-large"), num_classes=5) 6 优化器 = tf.keras.optimizers.Adam(learning_rate=1e-5) 7 优化器 = tf.mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')

13 帧

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_utils.py 在 from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs） 第399章 400 --> 401 model(model.dummy_inputs, training=False) # 使用虚拟输入构建网络 402 403 assert os.path.isfile(resolved_archive_file), "Error retrieving file {}".format(resolved_archive_file)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在 调用(self, *args, **kwargs) 966与base_layer_utils.autocast_context_manager（ 第967章 --> 968 个输出 = self.call(cast_inputs, *args, **kwargs) 第969章 970 self._set_mask_metadata（输入，输出，输入掩码）

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_roberta.py 通话中（自我，输入，**kwargs） 222 第223章 --> 224 个输出 = self.roberta(inputs, **kwargs) 225个返回输出 226

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在 调用(self, *args, **kwargs) 966与base_layer_utils.autocast_context_manager（ 第967章 --> 968 个输出 = self.call(cast_inputs, *args, **kwargs) 第969章 970 self._set_mask_metadata（输入，输出，输入掩码）

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py 在通话中（自我，输入，attention_mask，token_type_ids，position_ids， head_mask、inputs_embeds、培训） 第567章 568 --> 569 embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training) 570 编码器输出 = self.encoder（[嵌入输出，扩展注意力掩码，头掩码]，训练=训练） 第571章

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在 调用(self, *args, **kwargs) 966与base_layer_utils.autocast_context_manager（ 第967章 --> 968 个输出 = self.call(cast_inputs, *args, **kwargs) 第969章 970 self._set_mask_metadata（输入，输出，输入掩码）

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py 通话中（自我，输入，模式，训练） 第146章 147如果模式==“嵌入”： --> 148 返回 self._embedding(inputs, training=training) 149 elif 模式 == “线性”： 150 返回 self._linear(inputs)

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_roberta.py 在_embedding（自我，输入，培训） 79 position_ids = self.create_position_ids_from_inputs_embeds（inputs_embeds） 80 ---> 81 return super()._embedding([input_ids, position_ids, token_type_ids, inputs_embeds], training=training) 82 83

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py 在_embedding（自我，输入，培训） 173 174 个嵌入 = inputs_embeds + position_embeddings + token_type_embeddings --> 175 个嵌入 = self.LayerNorm(embeddings) 176 个嵌入 = self.dropout（嵌入，训练 = 训练） 177 个返回嵌入

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在 调用(self, *args, **kwargs) 962 # 渴望在数据张量上执行。 963 与 backend.name_scope(self._name_scope()): --> 964 self._maybe_build(输入) 第965章 966 with base_layer_utils.autocast_context_manager(

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在_maybe_build（自我，输入）2406 self._dtype_policy = policy.Policy(dtype) 2407 input_shapes = 无 -> 2408 if all(hasattr(x, 'shape') for x in input_list): 2409 input_shapes = nest.map_structure(lambda x: x.shape, inputs) 2410

只有在用户手动重写了构建方法时才调用build。

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py 在 (.0) 2406 self._dtype_policy = policy.Policy(dtype) 2407 input_shapes = 无 -> 2408 if all(hasattr(x, 'shape') for x in input_list): 2409 input_shapes = nest.map_structure(lambda x: x.shape, inputs) 2410

只有在用户手动重写了构建方法时才调用build。

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py 在形状（自我）1065 self._tensor_shape = tensor_shape.TensorShape(self._shape_tuple()) 1066 除了 core._NotOkStatusException 为 e： -> 1067 6.raise_from(core._status_to_exception(e.code, e.message), None) 1068 1069 return self._tensor_shape

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError：无法计算 AddV2，因为输入 #1（从零开始）是 预计是一个 bfloat16 张量，但它是一个浮点张量

我想我必须对模型进行一些调整？我怎样才能做到这一点 ？ 我正在使用 tensorflow 2.1 和 TPU v2。 我已经看到了这个线程，但我想它是使用 tensorflow 1.X 的，因为代码对我不起作用。 Memory reduction Tensorflow TPU v2/v3 bfloat16

Answer 1

我认为问题在于您正在尝试将使用全浮点数训练的预训练模型加载到 b16float 模型中。我认为这行不通。你必须从头开始训练。