У меня есть проблема с 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 (последний вызов последним)

in () 3 с помощью strategy.scope (): 4 ----> 5 model = CustomModel (TFXLMRobertaModel.from_pretrained ("jplu / tf-xlm-roberta-large"), num_classes = 5) 6 optimizer = tf.keras. optimizers.Adam (Learning_rate = 1e-5) 7 optimizer = 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 вернуть load_pytorch_checkpoint_in_tf2_model_fileted_arch, разрешить 401 model (model.dummy_inputs, training = False) # построить сеть с фиктивными входами 402 403 assert os.path.isfile (resolved_archive_file), "Ошибка при получении файла {}". Формат (resolved_archive_file)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в call (self, * args, ** kwargs) 966 с base_layer_utils.autocast_context_manager (967 self._compute_dtype): -> 968 output = self.call (cast_inputs, * args, ** kwargs) 969 self._handle_activity_regularization (входы, выходы) 970 self._set_mask_metadata (входы, выходы, input_masks)

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_roberta.py in call (self, inputs, ** kwargs) 222 223 "" "-> 224 output = self.roberta (inputs, ** kwargs) 225 обратных выходов 226

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в call (self, * args, ** kwargs) 966 с base_layer_utils.autocast_context_manager (967 self._compute_dtype): -> 968 output = self.call (cast_inputs, * args, ** kwargs) 969 self._handle_activity_regularization (входы, выходы) 970 self._set_mask_metadata (входы, выходы, input_masks)

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py в вызове (self, input, Внимание_mask, token_type_ids, position_ids, head_mask, inputs_embeds, training) 567 # head_mask = tf.constant ([0] * self.num_hidden_layers) 568 -> 569 embedding_output = self.embeddings ([input_ids, position_ids, token_type_ids, inputs_embeds], training = training) 570 encoder_outputs = self.encoder ([embedding_output, extended_attention_mask, head_mask], training = training = training

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в call (self, * args, ** kwargs) 966 с base_layer_utils.autocast_context_manager (967 self._compute_dtype): -> 968 output = self.call (cast_inputs, * args, ** kwargs) 969 self._handle_activity_regularization (входы, выходы) 970 self._set_mask_metadata (входы, выходы, input_masks)

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py в вызове (self, input, mode, training) 146 "" "147 if mode ==" embedding ": -> 148 return self ._embedding (входы, обучение = обучение) 149 elif mode == "linear": 150 return self._linear (inputs)

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_roberta.py в _embedding (self, input, training) 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 = обучение) 82 83

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_bert.py в _embedding (собственное, входные данные, обучение) 173 174 встраивания = input_embeds + position_embeddings + token_type_embeddings -> 175 встраиваний = self.LayerNorm (встраивания) 176 вложений = self.dropout (вложения, обучение = обучение) 177 вложений возврата

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в call (self, * args, ** kwargs) 962 # Активное выполнение на тензоры данных. 963 с backend.name_scope (self._name_scope ()): -> 964 self._maybe_build (входы) 965 cast_inputs = self._maybe_cast_inputs (входы) 966 с base_layer_utils.autocast_context_manager (

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в _maybe_build (self, inputs) 2406 self._dtype_policy = policy.Policy (dtype) 2407 input_shapes = None -> 2408, если все (hasattr (x, 'shape') для x в 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 = None -> 2408, если all (hasattr (x, 'shape') для x в 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 in shape (self) 1065 self._tensor_shape = tensor_shape.TensorShape (self._shape_tuple ()) 1066, кроме core._NotOkStatusException как e: -> 1067 six.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 в raise_from (значение, from_value)

InvalidArgumentError: невозможно вычислить AddV2, поскольку вход №1 (с нулевым основанием) должен был быть тензором bfloat16, но является тензором с плавающей запятой

Полагаю, мне нужно что-то отлить по поводу модели? Как я могу это сделать ? Я использую tensorflow 2.1 и TPU v2. Я видел эту ветку, но это было с tenorflow 1.X, я полагаю, поскольку код не работает для меня. Уменьшение памяти Tensorflow TPU v2 / v3 bfloat16