Пример #1
0
  def _testExtendStep(self, sess, dec, encoder_outputs, tgts, num_hyps):
    p = self._DecoderParams()

    # Infer true source encoder length from the padding.
    src_enc_len = tf.reduce_sum(1 - encoder_outputs.padding, axis=0)
    src_enc_len = dec._ExpandToNumHyps(src_enc_len, num_hyps)

    # Run Fprop
    fprop_out = dec._FProp(dec.theta, encoder_outputs, tgts)
    l_out1 = fprop_out.softmax_input
    attention_map_fprop = fprop_out.attention

    # run ExtendStep
    prefix_states = py_utils.NestedMap()
    for i in range(6):
      layer_i_states = py_utils.NestedMap()
      # The first dim is for the decode step (sequence length).
      # Here's 0 as placeholder
      layer_i_states.key = tf.zeros([0, self.tgt_batch, p.model_dim])
      layer_i_states.value = tf.zeros([0, self.tgt_batch, p.model_dim])
      prefix_states['layer_%i' % i] = layer_i_states

    l_out2 = []
    per_step_atten_probs = []
    for i in range(5):
      l_i_out, prefix_states, atten_probs = dec.ExtendStep(
          dec.theta, encoder_outputs, tgts.ids[:, i], i, prefix_states)
      l_out2.append(l_i_out)
      per_step_atten_probs.append(atten_probs)
    l_out2 = tf.stack(l_out2)
    bs_atten_probs = tf.stack(per_step_atten_probs)

    attention_map_bs = py_utils.NestedMap(probs=bs_atten_probs)

    def _TransposeAttentions(x):
      return tf.transpose(x, [1, 0, 2])

    attention_map_bs = attention_map_bs.Transform(_TransposeAttentions)

    tf.global_variables_initializer().run()

    l_out1_v, l_out2_v, attention_map_fprop_v, attention_map_bs_v, src_enc_len_v = sess.run(
        [l_out1, l_out2, attention_map_fprop, attention_map_bs, src_enc_len])

    # Ensure that FProp and BeamSearch output are the same.
    self.assertAllClose(l_out1_v, l_out2_v, rtol=1e-05, atol=1e-05)

    # Ensure that FProp and BeamSearch attention matrix is the same.
    self.assertAllClose(attention_map_fprop_v.probs, attention_map_bs_v.probs)

    print('attention map', attention_map_fprop_v.probs)

    # End-to-end test attention probs -- ensure EOS symbol and positions
    # behind EOS have 0 probability.
    for i in range(0, len(src_enc_len_v)):
      pos = int(src_enc_len_v[i]) - 1
      self.assertEqual(
          np.count_nonzero(attention_map_fprop_v.probs[i][:, pos:]), 0)
Пример #2
0
 def _TransformerSingleSourceInputs(self, depth=3, dtype=tf.float32):
   np.random.seed(NUMPY_RANDOM_SEED)
   source_vecs = tf.stack(
       [tf.constant(np.random.rand(2, depth), dtype=dtype) for _ in range(5)])
   source_padding = tf.transpose(
       tf.constant([[0, 0, 1, 1, 0], [1, 0, 0, 0, 1]], dtype=dtype))
   aux_source_vecs = tf.stack(
       [tf.constant(np.random.rand(2, depth), dtype=dtype) for _ in range(7)])
   aux_source_paddings = tf.transpose(
       tf.constant(
           [[0, 1, 0, 1, 0, 1, 0], [1, 0, 1, 0, 1, 0, 1]], dtype=dtype))
   return source_vecs, source_padding, aux_source_vecs, aux_source_paddings
Пример #3
0
def BatchMakeRotationMatrix(yaw, clockwise=False):
    """Create a Nx3x3 rotation matrix from yaw.

  Args:
    yaw: float tensor representing a yaw angle in radians.
    clockwise: Whether to have the rotation be applied clockwise (True) or
      counter-clockwise (False). Defaults to counter-clockwise to maintain
      same semantics to MakeRotationMatrix.

  Returns:
    A [N, 3, 3] tensor corresponding to a rotation matrix.
  """

    if clockwise:
        yaw = -yaw

    cos = tf.cos(yaw)
    sin = tf.sin(yaw)
    zero = tf.zeros_like(cos)
    one = tf.ones_like(cos)

    rotation_matrix = tf.stack(
        [cos, -sin, zero, sin, cos, zero, zero, zero, one],
        axis=-1)  # pyformat: disable
    rotation_matrix = tf.reshape(rotation_matrix, [-1, 3, 3])

    return rotation_matrix
Пример #4
0
 def _process(record):
     num, = tf.py_func(str_to_num, [record], [tf.float32])
     num = tf.stack([num, tf.square(num)])
     if use_nested_map:
         return py_utils.NestedMap(record=record, num=num), 1
     else:
         return [record, num], 1
Пример #5
0
def SplitTensors(xs, num_splits):
  """Splits tensors in `xs` evenly into num_splits along the 1st dimenion.

  Args:
    xs: A tuple of tensors. Each tensor's 1st dimension is the same size.
    num_splits: A python integer.

  Returns:
    A tuple of lists of tensors, num elements in the tuple = len(xs).

    i-th element in each list corresponds to i-th split of each tensor in xs
    along the first dimension of each tensor.
  """
  # assert first dim of all tensors in xs is equal
  batch_dims = [tf.shape(x)[0] for x in xs]
  all_batch_dims = tf.stack(batch_dims)

  all_batch_dims = py_utils.with_dependencies([
      py_utils.assert_equal(
          all_batch_dims,
          tf.shape(xs[0])[0],
          message='first dim of tensors in xs must match'),
      py_utils.assert_greater_equal(
          tf.shape(xs[0])[0],
          num_splits,
          message='first dim of tensors in xs must be greater than num_splits')
  ], all_batch_dims)

  splits = ComputeSplits(tf.shape(xs[0])[0], num_splits)
  # add the above assertion into the compute graph
  splits = py_utils.with_dependencies([all_batch_dims], splits)
  split_xs = [tf.split(axis=0, num_or_size_splits=splits, value=x) for x in xs]

  return split_xs
def _EncodeRandomJpegs(sizes):
    images = [
        tf.cast(
            tf.random.uniform([height, width, 3], maxval=256, dtype=tf.int32),
            tf.uint8) for height, width in sizes
    ]
    return tf.stack([tf.io.encode_jpeg(image) for image in images])
Пример #7
0
def SequenceAppendToken(x, x_paddings, token, extend=False):
    """Appends <token> to sequence `x`.

  Args:
    x: A sequence of tokens of shape [batch_size, x_len_max].
    x_paddings: The paddings of `x`.
    token: The token to append (of type integer).
    extend: Whether to extend `x` along the length dimension, this must be true
      for any sequence length in `x` that is `x_len_max` or else an invalid
      sequence will be emitted.

  Returns:
    A tuple.
      - The new sequence, Tensor of shape [batch_size, x_len_max].
      - The new paddings, Tensor of shape [batch_size, x_len_max].
  """
    batch_size = py_utils.GetShape(x)[0]
    x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
    if extend:
        x = tf.pad(x, [[0, 0], [0, 1]])
    # Mask all invalid entries of `x` to 0.
    x *= tf.sequence_mask(x_len, py_utils.GetShape(x)[1], x.dtype)
    # Append the <token> based on `x_len`.
    x += tf.scatter_nd(tf.stack([tf.range(batch_size), x_len], axis=1),
                       tf.cast(tf.fill([batch_size], token), x.dtype),
                       py_utils.GetShape(x))
    x_paddings = 1 - tf.sequence_mask(x_len + 1,
                                      py_utils.GetShape(x)[1],
                                      x_paddings.dtype)
    return x, x_paddings
Пример #8
0
    def CreateDenseCoordinates(self, ranges):
        """Create a matrix of coordinate locations corresponding to a dense grid.

    Example: To create (x, y) coordinates corresponding over a 10x10 grid with
    step sizes 1, call ``CreateDenseCoordinates([(1, 10, 10), (1, 10, 10)])``.

    Args:
      ranges: A list of 3-tuples, each tuple is expected to contain (min, max,
        num_steps). Each list element corresponds to one dimesion. Each tuple
        will be passed into np.linspace to create the values for a single
        dimension.

    Returns:
      tf.float32 tensor of shape [total_points, len(ranges)], where
      total_points = product of all num_steps.

    """
        total_points = int(np.prod([r_steps for _, _, r_steps in ranges]))
        cycle_steps = total_points
        stack_coordinates = []

        for r_start, r_stop, r_steps in ranges:
            values = tf.lin_space(tf.cast(r_start, tf.float32),
                                  tf.cast(r_stop, tf.float32),
                                  tf.cast(r_steps, tf.int32))
            cycle_steps //= r_steps
            gather_idx = (tf.range(total_points) // cycle_steps) % r_steps
            stack_coordinates.append(tf.gather(values, gather_idx))

        return tf.stack(stack_coordinates, axis=1)
Пример #9
0
    def testStep(self):
        p = self._testParams(dtype=tf.float32)
        with self.session(use_gpu=True) as sess:
            lm = p.Instantiate()
            inputs, paddings, _ = self._testInputs(dtype=tf.float32,
                                                   last_padding=0.0)

            sess.run(tf.global_variables_initializer())
            xent_output, _ = lm.FPropDefaultTheta(inputs=inputs,
                                                  paddings=paddings)
            logits1 = xent_output.logits

            time, batch = 5, 3
            prefix_states = lm.zero_state(lm.theta, batch)
            logits2 = []
            for i in range(time):
                l_i_out, prefix_states = lm.Step(lm.theta, inputs[i, :, :],
                                                 paddings[i, :], prefix_states)
                logits2.append(l_i_out.logits)
            logits2 = tf.stack(logits2)

            tf.global_variables_initializer().run()
            logits1_v, logits2_v = sess.run([logits1, logits2])
            print('xformer logits1_v', logits1_v)
            print('xformer logits2_v', logits2_v)
            self.assertAllClose(logits1_v, logits2_v)
Пример #10
0
 def Step(recurrent_theta, state0, inputs):
   """Computes one decoder step."""
   del inputs
   with tf.name_scope('single_sampler_step'):
     # Compute logits and states.
     bs_result, bs_state1 = pre_step_callback(
         recurrent_theta.theta,
         recurrent_theta.encoder_outputs,
         tf.expand_dims(state0.ids, 1),  # [batch, 1].
         state0.bs_state,
         num_hyps_per_beam=1)
     batch = tf.shape(bs_result.log_probs)[0]
     state1 = py_utils.NestedMap(timestep=state0.timestep + 1)
     state1.logits = bs_result.log_probs
     # Sample ids from logits. [batch].
     state1.ids = tf.reshape(
         tf.random.stateless_multinomial(
             state1.logits / p.temperature,
             num_samples=1,
             seed=tf.stack([recurrent_theta.random_seed, state0.timestep]),
             output_dtype=state0.ids.dtype,
             name='sample_next_id'), [batch])
     if 'is_last_chunk' in bs_result and p.target_eoc_id >= 0:
       state1.ids = tf.where(
           tf.logical_and(bs_result.is_last_chunk,
                          tf.equal(state1.ids, p.target_eoc_id)),
           tf.fill(tf.shape(state1.ids), p.target_eos_id), state1.ids)
     state1.bs_state = post_step_callback(recurrent_theta.theta,
                                          recurrent_theta.encoder_outputs,
                                          state1.ids, bs_state1)
   return state1, py_utils.NestedMap()
Пример #11
0
 def _IgnoreZCoordinate(bboxes):
   """Set z center to 0, and z dimension to 1."""
   num_bboxes = py_utils.GetShape(bboxes, 1)[0]
   return tf.stack([
       bboxes[:, 0], bboxes[:, 1], tf.zeros((num_bboxes,)),
       bboxes[:, 3], bboxes[:, 4], tf.ones((num_bboxes,)),
       bboxes[:, 6]
   ], axis=1)   # pyformat: disable
Пример #12
0
 def _process(source_id, record):
   num, = tf.py_func(str_to_num, [record], [tf.float32])
   num = tf.stack([num, tf.square(num)])
   if use_nested_map:
     return py_utils.NestedMap(
         source_id=source_id, record=record, num=num), bucket_fn(num)
   else:
     return [source_id, record, num], bucket_fn(num)
Пример #13
0
 def Update(self, new_value):
   state0 = self.GetValue()
   state1 = tf.stack([
       state0[0] + new_value[0],
       tf.minimum(state0[1], new_value[1]),
       tf.maximum(state0[2], new_value[2]),
   ])
   self.SetValue(state1)
Пример #14
0
def _MaybeStackExtraTheta(theta, all_vars, repeat):
  var_set = set([key for key, _ in all_vars.FlattenItems()])
  values = []
  for key, value in theta.FlattenItems():
    if key not in var_set and value is not None:
      # Replicate non-variable theta by p.repeat times.
      value = tf.stack([value] * repeat)
    values.append(value)
  return theta.Pack(values)
Пример #15
0
 def _testInputs(self, depth=3, dtype=tf.float32):
     np.random.seed(505837249)
     source_vecs = tf.stack([
         tf.constant(np.random.rand(2, depth), dtype=dtype)
         for _ in range(5)
     ])
     source_padding = tf.constant([[0, 0, 1, 1, 0], [1, 0, 0, 0, 1]],
                                  dtype=dtype)
     aux_source_vecs = tf.stack([
         tf.constant(np.random.rand(2, depth), dtype=dtype)
         for _ in range(7)
     ])
     aux_source_paddings = tf.constant(
         [[0, 1, 0, 1, 0, 1, 0], [1, 0, 1, 0, 1, 0, 1]], dtype=dtype)
     source_padding = tf.transpose(source_padding)
     aux_source_paddings = tf.transpose(aux_source_paddings)
     return (source_vecs, source_padding, aux_source_vecs,
             aux_source_paddings)
Пример #16
0
    def _testElmanHelper(self, seqlen, use_grad, stop_fn=None):
        with self.session() as sess:
            tf.set_random_seed(342462)

            batch = 3
            dims = 4
            theta = py_utils.NestedMap()
            theta.w = self.Rand([2 * dims, dims])
            theta.b = self.Rand([dims])
            state0 = py_utils.NestedMap()
            state0.h = self.Rand([batch, dims])
            inputs = py_utils.NestedMap()
            inputs.x = self.Rand([seqlen, batch, dims])

            # Static unrolled.
            s = state0
            out = []
            for i in range(seqlen):
                inp = py_utils.NestedMap()
                inp.x = inputs.x[i, :]
                s, _ = self.Elman(theta, s, inp)
                out += [s.h]
                if stop_fn and stop_fn(i + 1, theta, s):
                    out += [
                        tf.zeros_like(out[-1]) for _ in range(seqlen - i - 1)
                    ]
                    break
            acc0, final0 = tf.stack(out), s.h
            loss0 = tf.reduce_sum(acc0) + tf.reduce_sum(final0)
            (dw0, db0, dh0,
             di0) = tf.gradients(loss0, [theta.w, theta.b, state0.h, inputs.x])

            # Uses the Recurrent() library.
            acc1, final1 = recurrent.Recurrent(
                theta=theta,
                state0=state0,
                inputs=inputs,
                cell_fn=self.Elman,
                cell_grad=self.ElmanGrad if use_grad else None,
                stop_fn=stop_fn)
            acc1, final1 = acc1.h, final1.h
            loss1 = tf.reduce_sum(acc1) + tf.reduce_sum(final1)
            (dw1, db1, dh1,
             di1) = tf.gradients(loss1, [theta.w, theta.b, state0.h, inputs.x])

            # Fetches a bunch of values and compare them.
            (acc0, acc1, final0, final1, dw0, dw1, db0, db1, dh0, dh1, di0,
             di1) = sess.run([
                 acc0, acc1, final0, final1, dw0, dw1, db0, db1, dh0, dh1, di0,
                 di1
             ])
            self.assertAllClose(acc0, acc1)
            self.assertAllClose(final0, final1)
            self.assertAllClose(dw0, dw1)
            self.assertAllClose(db0, db1)
            self.assertAllClose(dh0, dh1)
            self.assertAllClose(di0, di1)
Пример #17
0
  def _TransformerAttentionLayerInputs(self, input_dim=4, dtype=tf.float32):
    np.random.seed(6348575)
    query_vec = tf.transpose(
        tf.stack([
            tf.constant(np.random.rand(2, input_dim), dtype=dtype)
            for _ in range(5)
        ]), [1, 0, 2])
    paddings = tf.constant([[0, 0, 1, 1, 0], [1, 0, 0, 0, 1]], dtype=dtype)
    aux_vec = tf.transpose(
        tf.stack([
            tf.constant(np.random.rand(2, input_dim), dtype=dtype)
            for _ in range(7)
        ]), [1, 0, 2])
    aux_paddings = tf.constant([[0, 1, 0, 1, 0, 1, 0], [1, 0, 1, 0, 1, 0, 1]],
                               dtype=dtype)

    segment_mask = tf.zeros([2, 1, 5, 5])
    return query_vec, paddings, aux_vec, aux_paddings, segment_mask
Пример #18
0
  def _testInputs(self, depth=3, dtype=tf.float32):
    np.random.seed(505837249)
    source_vecs = tf.stack(
        [tf.constant(np.random.rand(2, depth), dtype=dtype) for _ in range(5)])
    source_padding = tf.constant([[0, 0, 1, 1, 0], [1, 0, 0, 0, 1]],
                                 dtype=dtype)
    aux_source_vecs = tf.stack(
        [tf.constant(np.random.rand(2, depth), dtype=dtype) for _ in range(7)])
    aux_source_paddings = tf.constant(
        [[0, 1, 0, 1, 0, 1, 0], [1, 0, 1, 0, 1, 0, 1]], dtype=dtype)
    source_padding = tf.transpose(source_padding)
    aux_source_paddings = tf.transpose(aux_source_paddings)

    input_task_arr = np.array([[0] * depth, [0] * depth])
    tgt_task_arr = np.array([[0] * depth] * 3)
    input_tasks = tf.constant(input_task_arr.tolist(), dtype=tf.int32)
    tgt_tasks = tf.constant(tgt_task_arr.tolist(), dtype=tf.int32)
    return (source_vecs, source_padding, aux_source_vecs, aux_source_paddings,
            input_tasks, tgt_tasks)
Пример #19
0
    def NMSIndices(self,
                   bboxes,
                   scores,
                   max_output_size,
                   nms_iou_threshold=0.3,
                   score_threshold=0.01):
        """Apply NMS to a series of 3d bounding boxes in 7-DOF format.

    Args:
      bboxes: A [num_boxes, 7] floating point Tensor of bounding boxes in [x, y,
        z, dx, dy, dz, phi] format.
      scores: A [num_boxes] floating point Tensor containing box
        scores.
      max_output_size: Maximum number of boxes to predict per input.
      nms_iou_threshold: IoU threshold to use when determining whether two boxes
        overlap for purposes of suppression.
      score_threshold: The score threshold passed to NMS that allows NMS to
        quickly ignore irrelevant boxes.

    Returns:
      The NMS indices and the mask of the padded indices.
    """
        bboxes = py_utils.HasShape(bboxes, [-1, 7])

        # Extract x, y, w, h, then convert to extrema.
        #
        # Note that we drop the rotation angle because we don't have an NMS
        # operation that takes rotation into account.
        bboxes_2d = tf.stack(
            [bboxes[:, 0], bboxes[:, 1], bboxes[:, 3], bboxes[:, 4]], axis=-1)
        bboxes_extrema = geometry.XYWHToBBoxes(bboxes_2d)

        # Compute NMS with padding; we use the padded version so this function can
        # be used in a map_fn.  This function returns the scalar number of boxes
        # for each example.
        #
        # We use an IoU threshold of 0.3 since our anchor boxes have rotations
        # that make the default IoU threshold of 0.5 possibly too high.
        nms_index_padded, num_valid = tf.image.non_max_suppression_padded(
            bboxes_extrema,
            scores,
            iou_threshold=nms_iou_threshold,
            max_output_size=max_output_size,
            score_threshold=score_threshold,
            pad_to_max_output_size=True)

        # Return the mask of valid indices instead of just a scalar number.
        mask = tf.concat(
            [tf.ones([num_valid]),
             tf.zeros([max_output_size - num_valid])],
            axis=0)

        nms_index_padded = tf.where(mask > 0, nms_index_padded,
                                    tf.zeros_like(nms_index_padded))
        return nms_index_padded, mask
Пример #20
0
 def _MakeTransformTestRotationMatrices(self, batch_size):
   # Make a batch of 4x4 transformation matrices that only has rotation around
   # the z-axis (world rotation).
   rot_matrices = []
   for _ in range(batch_size):
     rot_matrix = geometry._MakeRotationMatrix(tf.random_uniform([]), 0., 0.)
     # Embed rotation matrix into a 4 x 4 matrix
     rot_matrix = tf.pad(rot_matrix, [[0, 1], [0, 1]]) + tf.diag([0, 0, 0, 1.])
     rot_matrices.append(rot_matrix)
   transforms = tf.stack(rot_matrices, axis=0)
   return transforms
Пример #21
0
 def _MakeTransformTestTranslationMatrices(self, batch_size):
     # Make a batch of 4x4 transformation matrices that translate in all
     # directions.
     translation_matrices = []
     for _ in range(batch_size):
         translation_matrix = tf.random_uniform([3, 1])
         translation_matrix = tf.pad(translation_matrix, [[0, 1], [3, 0]])
         translation_matrix += tf.diag([1., 1., 1., 1.])
         translation_matrices.append(translation_matrix)
     transforms = tf.stack(translation_matrices, axis=0)
     return transforms
Пример #22
0
        def Step(recurrent_theta, state0, inputs):
            """Computes one decoder step."""
            if p.use_recurrent:
                del inputs
            with tf.name_scope('single_sampler_step'):
                # Compute logits and states.
                bs_result, bs_state1 = pre_step_callback(
                    decoder_theta,
                    recurrent_theta.encoder_outputs,
                    tf.expand_dims(state0.ids, 1),  # [batch, 1].
                    state0.bs_state,
                    num_hyps_per_beam=p.num_hyps_per_beam)
                batch = tf.shape(bs_result.log_probs)[0]
                state1 = py_utils.NestedMap(timestep=state0.timestep + 1)
                state1.logits = bs_result.log_probs

                if p.top_k > 0:
                    topk_logits, topk_ids = tf.math.top_k(state1.logits,
                                                          k=p.top_k)
                    sample_logits = tf.nn.log_softmax(
                        topk_logits) if p.top_k_renormalize else topk_logits
                else:
                    sample_logits = state1.logits

                # Sample ids from logits. [batch].
                ids = tf.reshape(
                    tf.random.stateless_categorical(
                        sample_logits / p.temperature,
                        num_samples=1,
                        seed=tf.stack(
                            [recurrent_theta.random_seed, state0.timestep]),
                        dtype=state0.ids.dtype,
                        name='sample_next_id'), [batch])
                state1.ids = tf.gather(topk_ids, ids, axis=1,
                                       batch_dims=1) if p.top_k > 0 else ids

                if 'is_last_chunk' in bs_result and p.target_eoc_id >= 0:
                    state1.ids = tf.where(
                        tf.math.logical_and(
                            bs_result.is_last_chunk,
                            tf.equal(state1.ids, p.target_eoc_id)),
                        tf.fill(tf.shape(state1.ids), p.target_eos_id),
                        state1.ids)
                state1.bs_state = post_step_callback(
                    decoder_theta, recurrent_theta.encoder_outputs, state1.ids,
                    bs_state1)
            if p.use_recurrent:
                return state1, py_utils.NestedMap()
            else:
                inputs.ids = inputs.ids.write(state0.timestep, state1.ids)
                inputs.logits = inputs.logits.write(state0.timestep,
                                                    state1.logits)
                return (recurrent_theta, state1, inputs)
Пример #23
0
def BBoxCorners(bboxes):
    """Extract the corner points from a 7-DOF bbox representation.

  Args:
    bboxes: A [batch, num_boxes, 7] floating point bounding box representation
      ([x, y, z, dx, dy, dz, phi]).

  Returns:
    A [batch, num_boxes, 8, 3] floating point Tensor containing
      the corner (x, y, z) points for every bounding box.
  """
    # Code adapted from vale/soapbox codebase.
    #
    # Corners in normalized box frame (unit cube centered at origin).
    #
    # Dimensions is [length, width, height].
    corners = tf.constant([
        [0.5, 0.5, 0.5],  # top
        [-0.5, 0.5, 0.5],  # top
        [-0.5, -0.5, 0.5],  # top
        [0.5, -0.5, 0.5],  # top
        [0.5, 0.5, -0.5],  # bottom
        [-0.5, 0.5, -0.5],  # bottom
        [-0.5, -0.5, -0.5],  # bottom
        [0.5, -0.5, -0.5],  # bottom
    ])

    batch, nb, _ = py_utils.GetShape(bboxes, 3)

    # Extract location, dimension, and rotation.
    location = bboxes[:, :, :3]
    dimensions = bboxes[:, :, 3:6]
    phi_world = bboxes[:, :, 6]

    # Convert rotation_phis into rotation matrices along unit z.
    cos = tf.cos(phi_world)
    sin = tf.sin(phi_world)
    zero = tf.zeros_like(cos)
    one = tf.ones_like(cos)
    rotations_world = tf.reshape(
        tf.stack([cos, -sin, zero, sin, cos, zero, zero, zero, one], axis=2),
        [batch, nb, 3, 3])

    # Create axis-aligned corners from length/width/height.
    corners = tf.einsum('bni,ji->bnji', dimensions, corners)

    # Rotate the corners coordinates to the rotated world frame.
    corners = tf.einsum('bnij,bnkj->bnki', rotations_world, corners)

    # Translate corners to the world location.
    corners = corners + tf.reshape(location, (batch, nb, 1, 3))
    return corners
Пример #24
0
    def testEvolvedTransformerDecoderLayerExtendStep(self):
        with self.session(use_gpu=True) as sess:
            np.random.seed(6348575)
            depth = 4
            p = GPipeEvolvedTransformerDecoderLayer.Params()
            p.name = 'gpipe_evolved_transformer_decoder'
            p.source_dim = depth
            p.has_aux_atten = True
            p.mask_self_atten = True
            p.tr_double_heads_atten_tpl.num_attention_heads = 2
            p.tr_atten_tpl.num_attention_heads = 2
            p.transformer_tpl.tr_atten_tpl.num_attention_heads = 2
            et_decoder = GPipeEvolvedTransformerDecoderLayer(p)

            (source_vecs, _, aux_vecs,
             aux_paddings) = self._testInputs(depth=depth)
            source_padding = tf.zeros([5, 2])

            h1 = et_decoder.FPropDefaultTheta(
                aux_vecs,
                aux_paddings,
                source_vecs,
                source_padding,
                None,
                None,
                None,
                None,
            )[2]
            h2 = []

            double_head_attention_states = py_utils.NestedMap(
                key=tf.zeros([0, 2, 4]), value=tf.zeros([0, 2, 4]))
            transformer_layer_states = py_utils.NestedMap(
                key=tf.zeros([0, 2, 4]), value=tf.zeros([0, 2, 4]))
            branched_convs_input = tf.zeros([0, 2, 4])

            prefix_states = py_utils.NestedMap(
                double_head_attention_states=double_head_attention_states,
                transformer_layer_states=transformer_layer_states,
                branched_convs_input=branched_convs_input)

            for i in range(5):
                h, _, prefix_states = et_decoder.ExtendStep(
                    et_decoder.theta, source_vecs[i, :, :], prefix_states,
                    aux_vecs, aux_paddings)
                h2.append(h)

            h2 = tf.stack(h2)

            tf.global_variables_initializer().run()
            h1_v, h2_v = sess.run([h1, h2])
            self.assertAllClose(h1_v, h2_v)
Пример #25
0
    def QuantizeTensors(self, t_name, ts, eval_only=False):
        p = self.params
        # Always straddle a real zero point.
        if self.do_eval:
            # At eval/inference time, use the memorized range.
            # Important: Don't capture these variables in training mode so as to
            # avoid extra/unnecessary captures.
            min_var = self._GetQStateVar(t_name, 'min')
            max_var = self._GetQStateVar(t_name, 'max')
            return [
                self._MaybeFakeQuant(t, min_var, max_var, num_bits=p.bits)
                for t in ts
            ]
        else:
            # At training time, use the batch calculated min/max.
            accumulator_name = self._GetAccumulatorNameForTensor(t_name)
            # Calculate min/max for all tensors.
            batch_min = 0.0
            batch_max = 0.0
            for t in ts:
                batch_min = tf.minimum(tf.reduce_min(t), batch_min)
                batch_max = tf.maximum(tf.reduce_max(t), batch_max)

            # New state.
            state1 = tf.stack([1.0, batch_min, batch_max])
            self.accumulators[accumulator_name].Update(state1)

            # Results.
            ts_out = []
            for i, t in enumerate(ts):
                if eval_only:
                    # If only quantizing at eval time, still record ranges as above
                    # but don't quantize.
                    quant_t = t
                else:
                    # If quantizing during training, skip quantization if it produces
                    # NANs. Sometimes early in the training process, things are unstable
                    # and ranges can produce numerical instability that makes it
                    # impossible to perform a fake_quant.
                    quant_t = self._MaybeFakeQuant(t,
                                                   batch_min,
                                                   batch_max,
                                                   num_bits=p.bits)
                    # TODO(laurenzo): Plumb quant_t_has_nans through state and report.
                    quant_t_has_nans = tf.math.is_nan(quant_t)
                    quant_t = tf.where(quant_t_has_nans, t, quant_t)
                ts_out.append(quant_t)
                summary_utils.histogram(
                    '%s/%s_%d' % (self._qvars_scope.name, t_name, i), t)
            return ts_out
Пример #26
0
def SequenceConcat(x, x_paddings, y, y_paddings, pad=0):
    """Concats sequence `x` with sequence `y`.

  This function is length aware (based off the paddings).

  Args:
    x: A sequence of tokens of shape [batch_size, x_len_max].
    x_paddings: The paddings of `x`.
    y: A sequence of tokens of shape [batch_size, y_len_max].
    y_paddings: The paddings of `y`.
    pad: The <pad> token to fill the concatenated sequence (of type integer).

  Returns:
    A tuple.
      - Concatenation of `x` and `y` of shape
        [batch_size, x_len_max + y_len_max].
      - Paddings of the concatenation of shape
        [batch_size, x_len_max + y_len_max].
  """
    # Get the length (w/ eos).
    x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
    y_len = tf.cast(tf.round(tf.reduce_sum(1 - y_paddings, 1)), tf.int32)

    batch_size = py_utils.GetShape(x)[0]
    y_len_max = py_utils.GetShape(y)[1]

    # Pad `x` with necessary <pad>.
    x = tf.concat([x, tf.fill(py_utils.GetShape(y), pad)], 1)
    # Replace all <pad> with 0.
    x = tf.where(tf.not_equal(x, pad), x, tf.fill(py_utils.GetShape(x), 0))

    # Compute the write indices of `y` in `xy`.
    indices = tf.stack([
        tf.tile(tf.expand_dims(tf.range(batch_size), 1), [1, y_len_max]),
        (tf.tile(tf.expand_dims(tf.range(y_len_max), 0), [batch_size, 1]) +
         tf.expand_dims(x_len, 1)),
    ], 2)

    xy = x + tf.scatter_nd(indices, y, py_utils.GetShape(x))

    # We need to remap all <pad> to `pad`.
    xy = tf.where(
        tf.less(tf.expand_dims(tf.range(py_utils.GetShape(xy)[1]), 0),
                tf.expand_dims(x_len + y_len, 1)), xy,
        tf.fill(py_utils.GetShape(xy), pad))
    xy_paddings = 1 - tf.sequence_mask(x_len + y_len,
                                       py_utils.GetShape(xy)[1],
                                       x_paddings.dtype)
    return xy, xy_paddings
Пример #27
0
def _global_seed_from_inputs(input_floats):
  """Generates a random seed tensor based on input floats and mode key.

  Args:
    input_floats: a set of float input tensors that are derived from the input
      data (for example, input tokens). The important thing is that these are
      usually different for each batch.

  Returns:
    A tensor of shape=[2] with integer seed tensors derived from the inputs.
  """
  timestamp = tf.math.floormod(
      tf.cast(1000.0 * tf.timestamp(), dtype=tf.int64), 10000000000)
  input_sum = tf.cast(tf.reduce_sum(tf.math.abs(input_floats)), dtype=tf.int64)
  return tf.stack([timestamp + input_sum, timestamp - input_sum], axis=-1)
Пример #28
0
    def _Merge(*xs):
      """Broadcast all dimensions except the last, and concat on last dim."""

      # Stack all shapes and take max on each dimension to get leading shape.
      leading_shape = tf.stack([tf.shape(x)[:-1] for x in xs])
      leading_shape = tf.reduce_max(leading_shape, axis=0)
      # Broadcast each x.
      broadcast_xs = []
      for x in xs:
        broadcast_shape = tf.concat([leading_shape, tf.shape(x)[-1:]], axis=0)
        broadcast_xs.append(tf.broadcast_to(x, broadcast_shape))

      # Concat on last dimension.
      concat_xs = tf.concat(broadcast_xs, axis=-1)

      return concat_xs
Пример #29
0
    def _check_paddings(self, paddings):
        with tf.name_scope('check_paddings'):
            unpacked_paddings = tf.unstack(paddings)

            non_decr = []
            for t in unpacked_paddings:
                non_d = tf.is_non_decreasing(t)
                non_decr.append(non_d)
            all_non_decr = tf.stack(non_decr)

            paddings = py_utils.with_dependencies([
                tf.assert_equal(tf.reduce_any(tf.equal(paddings, 0.0)),
                                True,
                                message='must have at least one zero value.'),
                tf.assert_equal(
                    all_non_decr, True, message='must be non-decreasing')
            ], paddings)
            return paddings
Пример #30
0
    def testTransformerLayerExtendStep(self):
        with self.session(use_gpu=True) as sess:
            depth = 4
            np.random.seed(6348575)
            p = GPipeTransformerLayer.Params()
            p.name = 'transformer'
            p.source_dim = depth
            p.has_aux_atten = True
            p.mask_self_atten = True
            p.tr_fflayer_tpl.hidden_dim = 7
            p.tr_atten_tpl.num_attention_heads = 2
            transformer = GPipeTransformerLayer(p)

            (source_vecs, _, aux_vecs, aux_paddings, input_tasks,
             tgt_tasks) = self._testInputs(depth=depth)
            source_padding = tf.zeros([5, 2])

            output1 = transformer.FPropDefaultTheta(aux_vecs, aux_paddings,
                                                    source_vecs,
                                                    source_padding, None, None,
                                                    input_tasks, tgt_tasks)
            h1 = output1[2]
            out_src_task, out_tgt_task = output1[-2], output1[-1]

            h2 = []
            cached_source_vecs = tf.zeros([0, 2, 4])
            cached_source_contexts = tf.zeros([0, 2, 4])
            prefix_states = py_utils.NestedMap(key=cached_source_vecs,
                                               value=cached_source_contexts)
            for i in range(5):
                h, _, prefix_states = transformer.ExtendStep(
                    transformer.theta, source_vecs[i, :, :], prefix_states,
                    aux_vecs, aux_paddings)
                h2.append(h)

            h2 = tf.stack(h2)

            tf.global_variables_initializer().run()
            h1_v, h2_v = sess.run([h1, h2])
            self.assertAllClose(h1_v, h2_v)
            self.assertAllClose(out_src_task, input_tasks)
            self.assertAllClose(out_tgt_task, tgt_tasks)
            self.assertAllClose(
                h1_v[2][1], [1.10429943, -1.64884555, 0.15726769, -0.00250494])