Python tensorflow.bfloat16() Examples

def preprocess_for_train(image_bytes, use_bfloat16, image_size=IMAGE_SIZE, interpolation='bicubic'):
    """Preprocesses the given image for evaluation.

    Args:
      image_bytes: `Tensor` representing an image binary of arbitrary size.
      use_bfloat16: `bool` for whether to use bfloat16.
      image_size: image size.
      interpolation: image interpolation method

    Returns:
      A preprocessed image `Tensor`.
    """
    resize_method = tf.image.ResizeMethod.BICUBIC if interpolation == 'bicubic' else tf.image.ResizeMethod.BILINEAR
    image = _decode_and_random_crop(image_bytes, image_size, resize_method)
    image = _flip(image)
    image = tf.reshape(image, [image_size, image_size, 3])
    image = tf.image.convert_image_dtype(
        image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
    return image 

def __call__(self, x, *args, **kwargs):
    # Preprocessing: apply layer normalization
    #casting back to float32
    x = tf.cast(x, tf.bfloat16)
    y = self.layer_norm(x)
    #y = tf.cast(y, tf.float32)

    # Get layer output
    y = self.layer(y, *args, **kwargs)

    # Postprocessing: apply dropout and residual connection
    if self.train:
      mlperf_log.transformer_print(
            key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
            value=self.postprocess_dropout)
      y = tf.nn.dropout(y, 1 - (1 - self.postprocess_dropout))
    return x + y 

def _custom_getter(self):
    if self.hparams.weight_dtype == "bfloat16":
      if self.hparams.optimizer != "Adafactor":
        raise NotImplementedError(
            "weight_dtype=bfloat16 only implemented with Adafactor optimizer")
      activation_dtype = tf.float32
      if self.hparams.activation_dtype == "bfloat16":
        activation_dtype = tf.bfloat16
      return quantization.EighthPowerEncoding().custom_getter(
          activation_dtype=activation_dtype)
    elif self.hparams.activation_dtype == "bfloat16":
      return quantization.bfloat16_activations_var_getter
    elif mixed_precision_is_enabled(hparams=self.hparams):
      return quantization.float16_activations_var_getter
    else:
      return None 

def linear(self, x):
    """Computes logits by running x through a linear layer.

    Args:
      x: A float32 tensor with shape [batch_size, length, hidden_size]
    Returns:
      float32 tensor with shape [batch_size, length, vocab_size].
    """
    #with tf.compat.v1.tpu.bfloat16_scope():
    with tf.compat.v1.name_scope("presoftmax_linear"):
      #x = tf.cast(x, tf.bfloat16)
      batch_size = tf.shape(input=x)[0]
      length = tf.shape(input=x)[1]

      x = tf.reshape(x, [-1, self.hidden_size])
      logits = tf.matmul(x, self.shared_weights, transpose_b=True)
      #logits = tf.cast(logits, tf.float32)

      return tf.reshape(logits, [batch_size, length, self.vocab_size]) 

def bfloat16_activations_var_getter(getter, *args, **kwargs):
  """A custom getter function for float32 parameters and bfloat16 activations.

  Args:
    getter: custom getter
    *args: arguments
    **kwargs: keyword arguments
  Returns:
    variables with the correct dtype.
  Raises:
    KeyError: if "dtype" is not provided as a kwarg.
  """
  requested_dtype = kwargs["dtype"]
  if requested_dtype == tf.bfloat16:
    kwargs["dtype"] = tf.float32
  var = getter(*args, **kwargs)
  # This if statement is needed to guard the cast, because batch norm
  # assigns directly to the return value of this custom getter. The cast
  # makes the return value not a variable so it cannot be assigned. Batch
  # norm variables are always in fp32 so this if statement is never
  # triggered for them.
  if var.dtype.base_dtype != requested_dtype:
    var = tf.cast(var, requested_dtype)
  return var 

def preprocess_for_train(image_bytes, use_bfloat16, image_size=IMAGE_SIZE):
  """Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  """
  image = _decode_and_random_crop(image_bytes, image_size)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.image.convert_image_dtype(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
  return image 

def preprocess_for_eval(image_bytes, use_bfloat16, image_size=IMAGE_SIZE):
  """Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  """
  image = _decode_and_center_crop(image_bytes, image_size)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.image.convert_image_dtype(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
  return image 

def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.

  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.to_bfloat16(x)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.to_bfloat16(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.to_bfloat16(x_sign) 

def dataset_parser(value):
    keys_to_features = {
        'image/encoded': tf.FixedLenFeature((), tf.string, ''),
        'image/format': tf.FixedLenFeature((), tf.string, 'jpeg'),
        'image/class/label': tf.FixedLenFeature([], tf.int64, -1)
    }

    parsed = tf.parse_single_example(value, keys_to_features)
    image_bytes = tf.reshape(parsed['image/encoded'], shape=[])

    # Preprocess the images.
    image = tf.image.decode_jpeg(image_bytes)
    image = tf.image.random_flip_left_right(image)
    image = tf.image.resize_images(image, [IMAGE_SIZE, IMAGE_SIZE])
    image = tf.image.convert_image_dtype(
      image, dtype=tf.bfloat16)

    # Subtract one so that labels are in [0, 1000).
    label = tf.cast(
        tf.reshape(parsed['image/class/label'], shape=[]), dtype=tf.int32) - 1

    return image, label 

def custom_getter(self, activation_dtype=tf.bfloat16):
    """A custom getter that uses the encoding for bfloat16 and float32 vars.

    When a bfloat16 or float32 variable is requsted, an encoded float16
    varaible is created, which is then decoded and cast to a bfloat16
    activation.

    Args:
      activation_dtype: a dtype to which to convert the decoded value.

    Returns:
      a function.
    """
    def getter_fn(getter, *args, **kwargs):
      requested_dtype = kwargs["dtype"]
      if requested_dtype in (tf.bfloat16, tf.float32):
        kwargs["dtype"] = tf.bfloat16
        kwargs["initializer"] = _EncodingInitializer(
            kwargs["initializer"], self)
        ret = self._decode_with_identity_gradient(getter(*args, **kwargs))
        return tf.cast(ret, activation_dtype)
      return getter(*args, **kwargs)
    return getter_fn 

def custom_getter(self, activation_dtype=tf.bfloat16):
    """A custom getter that uses the encoding for bfloat16 and float32 vars.

    When a bfloat16 or float32 variable is requsted, an encoded float16
    varaible is created, which is then decoded and cast to a bfloat16
    activation.

    Args:
      activation_dtype: a dtype to which to convert the decoded value.

    Returns:
      a function.
    """
    def getter_fn(getter, *args, **kwargs):
      requested_dtype = kwargs["dtype"]
      if requested_dtype in (tf.bfloat16, tf.float32):
        kwargs["dtype"] = tf.bfloat16
        kwargs["initializer"] = _EncodingInitializer(
            kwargs["initializer"], self)
        ret = self._decode_with_identity_gradient(getter(*args, **kwargs))
        return tf.cast(ret, activation_dtype)
      return getter(*args, **kwargs)
    return getter_fn 

def bfloat16_activations_var_getter(getter, *args, **kwargs):
  """A custom getter function for float32 parameters and bfloat16 activations.

  Args:
    getter: custom getter
    *args: arguments
    **kwargs: keyword arguments
  Returns:
    variables with the correct dtype.
  Raises:
    KeyError: if "dtype" is not provided as a kwarg.
  """
  requested_dtype = kwargs["dtype"]
  if requested_dtype == tf.bfloat16:
    kwargs["dtype"] = tf.float32
  var = getter(*args, **kwargs)
  # This if statement is needed to guard the cast, because batch norm
  # assigns directly to the return value of this custom getter. The cast
  # makes the return value not a variable so it cannot be assigned. Batch
  # norm variables are always in fp32 so this if statement is never
  # triggered for them.
  if var.dtype.base_dtype != requested_dtype:
    var = tf.cast(var, requested_dtype)
  return var 

def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.

  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.

  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.

  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.to_float(cand1)
  cand2_f = tf.to_float(cand2)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret 

def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.

  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.to_bfloat16(x)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.to_bfloat16(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.to_bfloat16(x_sign) 

def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.

  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.to_bfloat16(x)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.to_bfloat16(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.to_bfloat16(x_sign) 

def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.

  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.

  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.

  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.to_float(cand1)
  cand2_f = tf.to_float(cand2)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret 

def preprocess_image(image_bytes,
                     is_training=False,
                     use_bfloat16=False,
                     image_size=IMAGE_SIZE,
                     interpolation='bicubic'):
    """Preprocesses the given image.

    Args:
      image_bytes: `Tensor` representing an image binary of arbitrary size.
      is_training: `bool` for whether the preprocessing is for training.
      use_bfloat16: `bool` for whether to use bfloat16.
      image_size: image size.
      interpolation: image interpolation method

    Returns:
      A preprocessed image `Tensor` with value range of [0, 255].
    """
    if is_training:
        return preprocess_for_train(image_bytes, use_bfloat16, image_size, interpolation)
    else:
        return preprocess_for_eval(image_bytes, use_bfloat16, image_size, interpolation) 

def preprocess_for_eval(image_bytes, use_bfloat16, image_size=IMAGE_SIZE, interpolation='bicubic'):
    """Preprocesses the given image for evaluation.

    Args:
      image_bytes: `Tensor` representing an image binary of arbitrary size.
      use_bfloat16: `bool` for whether to use bfloat16.
      image_size: image size.
      interpolation: image interpolation method

    Returns:
      A preprocessed image `Tensor`.
    """
    resize_method = tf.image.ResizeMethod.BICUBIC if interpolation == 'bicubic' else tf.image.ResizeMethod.BILINEAR
    image = _decode_and_center_crop(image_bytes, image_size, resize_method)
    image = tf.reshape(image, [image_size, image_size, 3])
    image = tf.image.convert_image_dtype(
        image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
    return image 

def model_fn(self, features):
    with tf.variable_scope(tf.get_variable_scope(), use_resource=True) as vs:
      self._add_variable_scope("model_fn", vs)
      transformed_features = self.bottom(features)

      if self.hparams.activation_dtype == "bfloat16":
        for k, v in sorted(six.iteritems(transformed_features)):
          if v.dtype == tf.float32:
            transformed_features[k] = tf.cast(v, tf.bfloat16)

      with tf.variable_scope("body") as body_vs:
        self._add_variable_scope("body", body_vs)
        log_info("Building model body")
        body_out = self.body(transformed_features)
      output, losses = self._normalize_body_output(body_out)

      if "training" in losses:
        log_info("Skipping T2TModel top and loss because training loss "
                 "returned from body")
        logits = output
      else:
        logits = self.top(output, features)
        losses["training"] = 0.0
        if (self._hparams.mode != tf.estimator.ModeKeys.PREDICT and
            self._hparams.mode != "attack"):
          losses["training"] = self.loss(logits, features)

      return logits, losses 

def decode(self, x):
    """Decode bfloat16 to float32."""
    raise NotImplementedError("decode not implemented") 

def encode(self, x, noise):
    """Encode float32 to bfloat16.

    Args:
      x: a float32 Tensor
      noise: a float32 Tensor with values in [0, 1), broadcastable to shape(x)

    Returns:
      a bfloat16 Tensor
    """
    raise NotImplementedError("encode not implemented") 

def _loss_single(self, logits, target_modality, feature):
    # The current bfloat16 version still uses float32 for most parts of backward
    # propagation to keep model quality, so cast back before computing the loss
    # value.
    if not target_modality:
      log_warn(_no_problem_err("loss"))
      return (tf.constant(0., dtype=tf.float32),
              tf.constant(1., dtype=tf.float32))

    loss_num, loss_den = target_modality.loss(logits, feature)
    loss_num *= self._problem_hparams.loss_multiplier

    if hasattr(self.hparams, "problem") and hasattr(
        self.hparams.problem, "task_list"):
      loss_num, loss_den, summaries = multi_problem.aggregate_task_losses(
          self.hparams,
          self._problem_hparams,
          logits,
          target_modality,
          feature
      )

      for key, val in summaries:
        tf.summary.scalar(key, val)

    return loss_num, loss_den 

def mzip(x,y):
    if x.dtype== tf.bfloat16:
      x = r_cast(x)
      y = r_cast(y)
    return zip(x,y) 

def __init__(self,
               is_training,
               use_bfloat16,
               transpose_input,
               data_dir,
               image_size=224,
               num_parallel_calls=64,
               num_cores=8,
               prefetch_depth_auto_tune=False,
               cache=False):
    """Create an input from TFRecord files.

    Args:
      is_training: `bool` for whether the input is for training
      use_bfloat16: If True, use bfloat16 precision; else use float32.
      transpose_input: 'bool' for whether to use the double transpose trick
      data_dir: `str` for the directory of the training and validation data;
          if 'null' (the literal string 'null') or implicitly False
          then construct a null pipeline, consisting of empty images
          and blank labels.
      image_size: size of input images
      num_parallel_calls: concurrency level to use when reading data from disk.
      num_cores: Number of prefetch threads
      prefetch_depth_auto_tune: Auto-tuning prefetch depths in input pipeline
      cache: if true, fill the dataset by repeating from its cache
    """
    super(ImageNetInput, self).__init__(
        is_training=is_training,
        image_size=image_size,
        use_bfloat16=use_bfloat16,
        num_cores=num_cores,
        prefetch_depth_auto_tune=prefetch_depth_auto_tune,
        transpose_input=transpose_input)
    self.data_dir = data_dir
    if self.data_dir == 'null' or not self.data_dir:
      self.data_dir = None
    self.num_parallel_calls = num_parallel_calls
    self.cache = cache 

def get_sequence_output(self):
    """Gets final hidden layer of encoder.

    Returns:
      float Tensor of shape [batch_size, seq_length, hidden_size] corresponding
      to the final hidden of the transformer encoder.

      In bfloat16 enabled execution, with only model covered in bfloat16 scope,
      return the output in float32. Other cases return as is
    """
    if self.bf16_scope == True:
      return  tf.cast(self.sequence_output, tf.float32)
    else :
      return self.sequence_output 

def get_pooled_output(self):
    """ 
      In bfloat16 enabled execution, with only model covered in bfloat16 scope,
      return the output in float32. Other cases return as is
    """
    if self.bf16_scope == True:
      return  tf.cast(self.pooled_output, tf.float32)
    else :
      return self.pooled_output 

def set_global_precision(dt):
  # Set Keras API precision
  global _keras_policy
  if dt == tf.bfloat16:
    _keras_policy=tf.keras.mixed_precision.experimental.Policy("mixed_bfloat16")

  # Set basic API precision
  set_rprecision(dt) 

def mzip(x,y):
    if x.dtype== tf.bfloat16:
      x = r_cast(x)
      y = r_cast(y)
    return zip(x,y) 

def get_sequence_output(self):
    """Gets final hidden layer of encoder.

    Returns:
      float Tensor of shape [batch_size, seq_length, hidden_size] corresponding
      to the final hidden of the transformer encoder.

      In bfloat16 enabled execution, with only model covered in bfloat16 scope,
      return the output in float32. Other cases return as is
    """
    if self.bf16_scope == True:
      return  tf.cast(self.sequence_output, tf.float32)
    else :
      return self.sequence_output 

def get_pooled_output(self):
    """ 
      In bfloat16 enabled execution, with only model covered in bfloat16 scope,
      return the output in float32. Other cases return as is
    """
    if self.bf16_scope == True:
      return  tf.cast(self.pooled_output, tf.float32)
    else :
      return self.pooled_output