Python v1.as_dtype方法代码示例

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def get_variable_dtype(
    master_dtype=tf.bfloat16,
    slice_dtype=tf.float32,
    activation_dtype=tf.float32):
  """Datatypes to use for the run.

  Args:
    master_dtype: string, datatype for checkpoints
      keep this the same between training and eval/inference
    slice_dtype: string, datatype for variables in memory
      must be tf.float32 for training
    activation_dtype: string, datatype for activations
      less memory usage if tf.bfloat16 but possible numerical issues
  Returns:
    a mtf.VariableDtype
  """
  return mtf.VariableDType(
      master_dtype=tf.as_dtype(master_dtype),
      slice_dtype=tf.as_dtype(slice_dtype),
      activation_dtype=tf.as_dtype(activation_dtype)) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def _get_tf_dtype(space):
    if isinstance(space, gym.spaces.Discrete):
      return tf.int32
    if isinstance(space, gym.spaces.Box):
      return tf.as_dtype(space.dtype)
    raise NotImplementedError() 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def master_dtype(self):
    return tf.as_dtype(self._hparams.master_dtype) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def slice_dtype(self):
    return tf.as_dtype(self._hparams.slice_dtype) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def activation_dtype(self):
    return tf.as_dtype(self._hparams.activation_dtype) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def variable_dtype(self):
    return mtf.VariableDType(
        tf.as_dtype(self._hparams.master_dtype),
        tf.as_dtype(self._hparams.slice_dtype),
        tf.as_dtype(self._hparams.activation_dtype)) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def to_spec(cls, instance):
    if isinstance(instance, tf.Tensor):
      return ExtendedTensorSpec.from_tensor(instance)
    elif isinstance(instance, TSPEC):
      return ExtendedTensorSpec.from_spec(instance)
    elif isinstance(instance, np.ndarray):
      return ExtendedTensorSpec(
          shape=instance.shape, dtype=tf.as_dtype(instance.dtype),
          is_extracted=True)
    raise ValueError(
        'We cannot convert {} with type {} to ExtendedTensorSpec'.format(
            instance, type(instance))) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def toy_model(features, mesh):
  """A toy model implemented by mesh tensorlfow."""
  batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
  io_dim = mtf.Dimension('io', FLAGS.io_size)

  master_dtype = tf.as_dtype(FLAGS.master_dtype)
  slice_dtype = tf.as_dtype(FLAGS.slice_dtype)
  activation_dtype = tf.as_dtype(FLAGS.activation_dtype)

  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
  x = mtf.cast(x, activation_dtype)
  h = x
  for lnum in range(1, FLAGS.num_hidden_layers + 2):
    if lnum + 1 == FLAGS.num_hidden_layers + 2:
      # output layer
      dim = io_dim
    elif lnum % 2 == 0:
      dim = mtf.Dimension('hidden_even', FLAGS.hidden_size)
    else:
      dim = mtf.Dimension('hidden_odd', FLAGS.hidden_size)
    h = mtf.layers.dense(
        h, dim,
        use_bias=False,
        master_dtype=master_dtype,
        slice_dtype=slice_dtype,
        name='layer_%d' % lnum)
  y = h
  g = tf.train.get_global_step()
  if FLAGS.step_with_nan >= 0:
    # Trigger NaN in the forward pass, this is used for testing whether
    # MeshTensorFlow can handle occasional NaN value.
    y += mtf.import_tf_tensor(
        mesh,
        tf.divide(
            0.0,
            tf.cond(tf.equal(g, FLAGS.step_with_nan), lambda: 0., lambda: 1.)),
        mtf.Shape([]))

  loss = mtf.reduce_mean(mtf.square(y - x))
  return y, loss 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def onnx2tf(dtype):
  return tf.as_dtype(mapping.TENSOR_TYPE_TO_NP_TYPE[_onnx_dtype(dtype)]) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def irdft_matrix(shape, dtype=tf.float32):
  """Matrix for implementing kernel reparameterization with `tf.matmul`.

  This can be used to represent a kernel with the provided shape in the RDFT
  domain.

  Example code for kernel creation, assuming 2D kernels:

  ```
  def create_kernel(init):
    shape = init.shape.as_list()
    matrix = irdft_matrix(shape[:2])
    init = tf.reshape(init, (shape[0] * shape[1], shape[2] * shape[3]))
    init = tf.matmul(tf.transpose(matrix), init)
    kernel = tf.Variable(init)
    kernel = tf.matmul(matrix, kernel)
    kernel = tf.reshape(kernel, shape)
    return kernel
  ```

  Args:
    shape: Iterable of integers. Shape of kernel to apply this matrix to.
    dtype: `dtype` of returned matrix.

  Returns:
    `Tensor` of shape `(prod(shape), prod(shape))` and dtype `dtype`.
  """
  shape = tuple(int(s) for s in shape)
  dtype = tf.as_dtype(dtype)
  size = np.prod(shape)
  rank = len(shape)
  matrix = np.identity(size, dtype=np.float64).reshape((size,) + shape)
  for axis in range(rank):
    matrix = fftpack.rfft(matrix, axis=axis + 1)
    slices = (rank + 1) * [slice(None)]
    if shape[axis] % 2 == 1:
      slices[axis + 1] = slice(1, None)
    else:
      slices[axis + 1] = slice(1, -1)
    matrix[tuple(slices)] *= np.sqrt(2)
  matrix /= np.sqrt(size)
  matrix = np.reshape(matrix, (size, size))
  return tf.constant(
      matrix, dtype=dtype, name="irdft_" + "x".join([str(s) for s in shape])) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import as_dtype [as 别名]
def __init__(self,
               symbols_to_logits_fn,
               vocab_size,
               batch_size,
               beam_size,
               alpha,
               max_decode_length,
               eos_id,
               padded_decode,
               dtype=tf.float32):
    """Initialize sequence beam search.

    Args:
      symbols_to_logits_fn: A function to provide logits, which is the
        interface to the Transformer model. The passed in arguments are:
          ids -> A tensor with shape [batch_size * beam_size, index].
          index -> A scalar.
          cache -> A nested dictionary of tensors [batch_size * beam_size, ...].
        The function must return a tuple of logits and the updated cache:
          logits -> A tensor with shape [batch * beam_size, vocab_size].
          updated cache -> A nested dictionary with the same structure as the
            input cache.
      vocab_size: An integer, the size of the vocabulary, used for topk
        computation.
      batch_size: An integer, the decode batch size.
      beam_size: An integer, number of beams for beam search.
      alpha: A float, defining the strength of length normalization.
      max_decode_length: An integer, the maximum number of steps to decode
        a sequence.
      eos_id: An integer. ID of end of sentence token.
      padded_decode: A bool, indicating if max_sequence_length padding is used
        for beam search.
      dtype: A tensorflow data type used for score computation. The default is
        tf.float32.
    """
    self.symbols_to_logits_fn = symbols_to_logits_fn
    self.vocab_size = vocab_size
    self.batch_size = batch_size
    self.beam_size = beam_size
    self.alpha = alpha
    self.max_decode_length = max_decode_length
    self.eos_id = eos_id
    self.padded_decode = padded_decode
    self.dtype = tf.as_dtype(dtype)