Tensorflow 2.1 TPU v2 сокращает использование памяти с помощью bfloat16 - PullRequest
Новая
       36

Tensorflow 2.1 TPU v2 сокращает использование памяти с помощью bfloat16

0 голосов
/ 25 апреля 2020

У меня есть некоторые проблемы с 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 optimizer = tf.keras.optimizers .Adam (learning_rate = 1e-5) 7 оптимизатор = tf.mixed_precision.LossScaleOptimizer (оптимизатор, loss_scale = 'dynamici c')

13 кадров

/ usr / local / lib / python3 .6 / dist-packages / transformers / modelling_tf_utils.py в from_pretrained (cls, pretrained_model_name_or_path, * model_args, ** kwargs) 399 return load_pytorch_checkpoint_in_tf2_model (модель, resolved_archive_file, модель модели - 400)> allow_missing_keys_keys_keys_se_set_set_set_set_set_set_set_set_set_setup dummy_inputs, training = False) # построить сеть с использованием фиктивных входов 402 403 assert os.path.isfile (resolved_archive_file), "Ошибка при получении файла {}". format (resolved_archive_file)

/ usr / local / lib / python3 .6 / dist-packages / tenorflow / python / keras / engine / base_layer.py в вызов (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 (входы, выходы, входные_маски)

/ usr /local/lib/python3.6/dist-packages/transformers/modeling_tf_roberta.py в вызове (self, input, ** kwargs) 222 223 "" "-> 224 выходов = self.roberta (входов, ** kwargs) 225 возвращаемых выходов 226

/ usr / local / lib / python3 .6 / dist-packages / tenorflow / python / keras / engine / base_layer.py в call (self , * args, ** kwargs) 966 с base_layer_utils.autocast_context_manager (967 self._compute_dtype): -> 968 выходов = self.call (cast_inputs, * args, ** kwargs) 969 self._handle_activity_regularization (входы, выходы) 970 self ._set_mask_metadata (входы, выходы, входные_маски)

/ usr / local / lib / python3 .6 / dist -packages / transformers / models_tf_bert.py в вызове (self, input, Внимание_Маска, token_type_ids, position_ids, head_mask, input_embeds, обучение) 567 # head_mask = tf.constant ([0] * self.num_hidden_layers) 568 -> 569 embedding_output = self.embeddings ([input_ids, position_ids, token_type_ids, input_embeds], обучение = обучение) 570 encoder_outputs = self.encoder ([embedding_output, extended_attention_mask, head_mask], обучение = обучение) 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 self._compute_dtype ): -> 968 выходов = self.call (cast_inputs, * args, ** kwargs) 969 self._handle_activity_regularization (входы, выходы) 970 self._set_mask_metadata (входы, выходы, входные_маски)

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

/ usr / local / lib / python3. 6 / dist-packages / transformers / моделирование_tf_roberta.py в _embedding (self, input, training) 79 position_ids = self.create_position_ids_from_inputs_embeds (input_embeds) 80 ---> 81 return super () ._ embedding ([input_ids, position_ids, token_type_ids, input_embeds) ], training = training) 82 83

/ usr / local / lib / python3 .6 / dist-packages / transformers / modelling_tf_bert.py в _embedding (self, input, training) 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 / /python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py в _maybe_build (self, input) 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, input) 2410

Только вызов build, если пользователь вручную переопределил сборку method.

/ usr / local / lib / python3 .6 / dist-packages / tenorflow / python / keras / engine / base_layer.py в (.0) 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 / lo cal / lib / python3 .6 / dist-packages / tenorflow / python / framework / ops.py в форме (self) 1065 self._tensor_shape = tenor_shape.TensorShape (self._shape_tuple ()) 1066 за исключением core._NotOkStatusException как e : -> 1067 six.raise_from (core._status_to_exception (e.code, e.message), None) 1068 1069 вернуть self._tensor_shape

/ usr / local / lib / python3 .6 / dist- packages / six.py в повышение_значении (значение, из_значения)

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

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

1 Ответ

0 голосов
/ 28 апреля 2020

Я думаю, что проблема в том, что вы пытаетесь загрузить предварительно обученную модель, обученную с полными поплавками, в модель b16float. Я не думаю, что это сработает. Вы должны тренироваться с нуля.

...