Python numpy.take_along_axis方法代码示例

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testSortIndicesExecution(self):
        # only 1 chunk when axis = -1
        raw = np.random.rand(100, 10)
        x = tensor(raw, chunk_size=10)

        r = sort(x, return_index=True)

        sr, si = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(sr, np.take_along_axis(raw, si, axis=-1))

        x = tensor(raw, chunk_size=(22, 4))

        r = sort(x, return_index=True)

        sr, si = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(sr, np.take_along_axis(raw, si, axis=-1))

        raw = np.random.rand(100)

        x = tensor(raw, chunk_size=23)

        r = sort(x, axis=0, return_index=True)

        sr, si = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(sr, raw[si]) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testArgsort(self):
        # only 1 chunk when axis = -1
        raw = np.random.rand(100, 10)
        x = tensor(raw, chunk_size=10)

        xa = argsort(x)

        r = self.executor.execute_tensor(xa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw), np.take_along_axis(raw, r, axis=-1))

        x = tensor(raw, chunk_size=(22, 4))

        xa = argsort(x)

        r = self.executor.execute_tensor(xa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw), np.take_along_axis(raw, r, axis=-1))

        raw = np.random.rand(100)

        x = tensor(raw, chunk_size=23)

        xa = argsort(x, axis=0)

        r = self.executor.execute_tensor(xa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw, axis=0), raw[r]) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testRelativePositionalEmbeddingLayer(self):
    with self.session(use_gpu=False):
      radius = 3
      p = layers.RelativePositionalEmbeddingLayer.Params().Set(
          name='rel_position_emb', radius=radius, dim=4)
      layer = p.Instantiate()
      indices = np.array([-5, -2, 0, 1, 4], dtype=np.int32)
      pos_emb = layer.FPropDefaultTheta(tf.convert_to_tensor(indices))

      self.evaluate(tf.global_variables_initializer())
      actual_pos_emb, full_emb = self.evaluate([pos_emb, layer.vars.w])

      clipped_indices = np.vectorize(lambda x: max(-radius, min(radius, x)))(
          indices) + radius
      expected_output = np.take_along_axis(full_emb,
                                           np.expand_dims(clipped_indices, -1),
                                           0)
      print('expected_position_embs:', expected_output)
      print('actual_position_embs:', actual_pos_emb)
      self.assertAllClose(actual_pos_emb, expected_output) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
    rng = rng_factory()
    i_shape = onp.array(i_shape)
    if axis is None:
      i_shape = [onp.prod(i_shape, dtype=onp.int64)]
    else:
      # Test the case where the size of the axis doesn't necessarily broadcast.
      i_shape[axis] *= 3
      i_shape = list(i_shape)
    def args_maker():
      x = rng(x_shape, dtype)
      n = onp.prod(x_shape, dtype=onp.int32) if axis is None else x_shape[axis]
      i = rng(i_shape, onp.int32) % (2 * n - 1) - (n - 1)
      return x, i

    lnp_op = lambda x, i: lnp.take_along_axis(x, i, axis=axis)

    if hasattr(onp, "take_along_axis"):
      onp_op = lambda x, i: onp.take_along_axis(x, i, axis=axis)
      self._CheckAgainstNumpy(lnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True,
                          check_incomplete_shape=True) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def cummin(x_val, x_key):
    """Get the cumulative minimum of `x_val` when ranked according to `x_key`.

    Parameters
    ----------
    x_val : :class:`numpy:numpy.ndarray` of shape (n, d)
        The array to get the cumulative minimum of along axis 0.
    x_key : :class:`numpy:numpy.ndarray` of shape (n, d)
        The array for ranking elements as to what is the minimum.

    Returns
    -------
    c_min : :class:`numpy:numpy.ndarray` of shape (n, d)
        The cumulative minimum array.
    """
    assert x_val.shape == x_key.shape
    assert x_val.ndim == 2
    assert not np.any(np.isnan(x_key)), "cummin not defined for nan key"

    n, _ = x_val.shape

    xm = np.minimum.accumulate(x_key, axis=0)
    idx = np.maximum.accumulate((x_key <= xm) * np.arange(n)[:, None])
    c_min = np.take_along_axis(x_val, idx, axis=0)
    return c_min 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def test_emd():
    x1 = torch.rand(20, 8192, 3).cuda()
    x2 = torch.rand(20, 8192, 3).cuda()
    emd = emdModule()
    start_time = time.perf_counter()
    dis, assigment = emd(x1, x2, 0.05, 3000)
    print("Input_size: ", x1.shape)
    print("Runtime: %lfs" % (time.perf_counter() - start_time))
    print("EMD: %lf" % np.sqrt(dis.cpu()).mean())
    print("|set(assignment)|: %d" % assigment.unique().numel())
    assigment = assigment.cpu().numpy()
    assigment = np.expand_dims(assigment, -1)
    x2 = np.take_along_axis(x2, assigment, axis = 1)
    d = (x1 - x2) * (x1 - x2)
    print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean())

#test_emd() 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testPartitionIndicesExecution(self):
        # only 1 chunk when axis = -1
        raw = np.random.rand(100, 10)
        x = tensor(raw, chunk_size=10)

        kth = [2, 5, 9]
        r = partition(x, kth, return_index=True)

        pr, pi = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(pr, np.take_along_axis(raw, pi, axis=-1))
        np.testing.assert_array_equal(np.sort(raw)[:, kth], pr[:, kth])

        x = tensor(raw, chunk_size=(22, 4))

        r = partition(x, kth, return_index=True)

        pr, pi = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(pr, np.take_along_axis(raw, pi, axis=-1))
        np.testing.assert_array_equal(np.sort(raw)[:, kth], pr[:, kth])

        raw = np.random.rand(100)

        x = tensor(raw, chunk_size=23)

        r = partition(x, kth, axis=0, return_index=True)

        pr, pi = self.executor.execute_tensors(r)
        np.testing.assert_array_equal(pr, np.take_along_axis(raw, pi, axis=-1))
        np.testing.assert_array_equal(np.sort(raw)[kth], pr[kth]) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testArgpartitionExecution(self):
        # only 1 chunk when axis = -1
        raw = np.random.rand(100, 10)
        x = tensor(raw, chunk_size=10)

        kth = [6, 3, 8]
        pa = argpartition(x, kth)

        r = self.executor.execute_tensor(pa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw)[:, kth], np.take_along_axis(raw, r, axis=-1)[:, kth])

        x = tensor(raw, chunk_size=(22, 4))

        pa = argpartition(x, kth)

        r = self.executor.execute_tensor(pa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw)[:, kth], np.take_along_axis(raw, r, axis=-1)[:, kth])

        raw = np.random.rand(100)

        x = tensor(raw, chunk_size=23)

        pa = argpartition(x, kth, axis=0)

        r = self.executor.execute_tensor(pa, concat=True)[0]
        np.testing.assert_array_equal(np.sort(raw, axis=0)[kth], raw[r][kth]) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testTopkExecution(self):
        raw1, order1 = np.random.rand(5, 6, 7), None
        raw2 = np.empty((5, 6, 7), dtype=[('a', np.int32), ('b', np.float64)])
        raw2['a'] = np.random.randint(1000, size=(5, 6, 7), dtype=np.int32)
        raw2['b'] = np.random.rand(5, 6, 7)
        order2 = ['b', 'a']

        for raw, order in [(raw1, order1), (raw2, order2)]:
            for chunk_size in [7, 4]:
                a = tensor(raw, chunk_size=chunk_size)
                for axis in [0, 1, 2, None]:
                    size = raw.shape[axis] if axis is not None else raw.size
                    for largest in [True, False]:
                        for to_sort in [True, False]:
                            for parallel_kind in ['tree', 'psrs']:
                                for k in [2, size - 2, size, size + 2]:
                                    r = topk(a, k, axis=axis, largest=largest, sorted=to_sort,
                                             order=order, parallel_kind=parallel_kind)

                                    result = self.executor.execute_tensor(r, concat=True)[0]

                                    if not to_sort:
                                        result = self._handle_result(result, axis, largest, order)
                                    expected = self._topk_slow(raw, k, axis, largest, order)
                                    np.testing.assert_array_equal(result, expected)

                                    r = topk(a, k, axis=axis, largest=largest,
                                             sorted=to_sort, order=order,
                                             parallel_kind=parallel_kind,
                                             return_index=True)

                                    ta, ti = self.executor.execute_tensors(r)
                                    raw2 = raw
                                    if axis is None:
                                        raw2 = raw.flatten()
                                    np.testing.assert_array_equal(ta, np.take_along_axis(raw2, ti, axis))
                                    if not to_sort:
                                        ta = self._handle_result(ta, axis, largest, order)
                                    np.testing.assert_array_equal(ta, expected) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def topk(input, k, dim, descending=True):
    topk_indices = argtopk(input, k, dim, descending)
    return np.take_along_axis(input, topk_indices, axis=dim) 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def _gather_feat(feat, ind, mask=None):
    dim  = feat.shape[2]
    ind = np.repeat(ind[:, :, np.newaxis], dim, axis=2)
    feat = np.take_along_axis(feat, ind, 1) 
    if mask is not None:
        mask = np.expand_dims(mask, 2).reshape(feat.shape)
        feat = feat[mask]
        feat = feat.reshape(-1, dim)
    return feat 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def call(self, x, training=None):
        import numpy as np

        x_ = x.numpy()
        if not is_channels_first(self.data_format):
            x_ = x_.transpose((0, 3, 1, 2))

        heatmap = x_[:, :-4]
        wh = x_[:, -4:-2]
        reg = x_[:, -2:]
        batch, _, out_h, out_w = heatmap.shape

        heatmap_flat = heatmap.reshape((batch, -1))
        indices = np.argsort(heatmap_flat)[:, -self.topk:]
        scores = np.take_along_axis(heatmap_flat, indices=indices, axis=-1)
        topk_classes = (indices // (out_h * out_w)).astype(dtype=np.float32)
        topk_indices = indices % (out_h * out_w)
        topk_ys = (topk_indices // out_w).astype(dtype=np.float32)
        topk_xs = (topk_indices % out_w).astype(dtype=np.float32)
        center = reg.transpose((0, 2, 3, 1)).reshape((batch, -1, 2))
        wh = wh.transpose((0, 2, 3, 1)).reshape((batch, -1, 2))
        xs = np.take_along_axis(center[:, :, 0], indices=topk_indices, axis=-1)
        ys = np.take_along_axis(center[:, :, 1], indices=topk_indices, axis=-1)
        topk_xs = topk_xs + xs
        topk_ys = topk_ys + ys
        w = np.take_along_axis(wh[:, :, 0], indices=topk_indices, axis=-1)
        h = np.take_along_axis(wh[:, :, 1], indices=topk_indices, axis=-1)
        half_w = 0.5 * w
        half_h = 0.5 * h
        bboxes = tf.stack((topk_xs - half_w, topk_ys - half_h, topk_xs + half_w, topk_ys + half_h), axis=-1)

        bboxes = bboxes * self.scale
        topk_classes = tf.expand_dims(topk_classes, axis=-1)
        scores = tf.expand_dims(scores, axis=-1)
        result = tf.concat((bboxes, topk_classes, scores), axis=-1)
        return result 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def __call__(self, x):
        import numpy as np

        heatmap = x[:, :-4].array
        wh = x[:, -4:-2].array
        reg = x[:, -2:].array
        batch, _, out_h, out_w = heatmap.shape

        heatmap_flat = heatmap.reshape((batch, -1))
        indices = np.argsort(heatmap_flat)[:, -self.topk:]
        scores = np.take_along_axis(heatmap_flat, indices=indices, axis=-1)
        topk_classes = (indices // (out_h * out_w)).astype(dtype=np.float32)
        topk_indices = indices % (out_h * out_w)
        topk_ys = (topk_indices // out_w).astype(dtype=np.float32)
        topk_xs = (topk_indices % out_w).astype(dtype=np.float32)
        center = reg.transpose((0, 2, 3, 1)).reshape((batch, -1, 2))
        wh = wh.transpose((0, 2, 3, 1)).reshape((batch, -1, 2))
        xs = np.take_along_axis(center[:, :, 0], indices=topk_indices, axis=-1)
        ys = np.take_along_axis(center[:, :, 1], indices=topk_indices, axis=-1)
        topk_xs = topk_xs + xs
        topk_ys = topk_ys + ys
        w = np.take_along_axis(wh[:, :, 0], indices=topk_indices, axis=-1)
        h = np.take_along_axis(wh[:, :, 1], indices=topk_indices, axis=-1)
        half_w = 0.5 * w
        half_h = 0.5 * h
        bboxes = F.stack((topk_xs - half_w, topk_ys - half_h, topk_xs + half_w, topk_ys + half_h), axis=-1)

        bboxes = bboxes * self.scale
        topk_classes = F.expand_dims(topk_classes, axis=-1)
        scores = F.expand_dims(scores, axis=-1)
        result = F.concat((bboxes, topk_classes, scores), axis=-1)
        return result 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def __call__(self, x):
        raw_pre_features = self.backbone(x)

        rpn_score = self.navigator_unit(raw_pre_features)
        rpn_score.to_cpu()
        all_cdds = [np.concatenate((y.reshape(-1, 1), self.edge_anchors.copy()), axis=1)
                    for y in rpn_score.array]
        top_n_cdds = [hard_nms(y, top_n=self.top_n, iou_thresh=0.25) for y in all_cdds]
        top_n_cdds = np.array(top_n_cdds)
        top_n_index = top_n_cdds[:, :, -1].astype(np.int64)
        top_n_index = np.array(top_n_index, dtype=np.int64)
        top_n_prob = np.take_along_axis(rpn_score.array, top_n_index, axis=1)

        batch = x.shape[0]
        x_pad = F.pad(x, pad_width=self.pad_width, mode="constant", constant_values=0)
        part_imgs = []
        for i in range(batch):
            for j in range(self.top_n):
                y0, x0, y1, x1 = tuple(top_n_cdds[i][j, 1:5].astype(np.int64))
                x_res = F.resize_images(
                    x_pad[i:i + 1, :, y0:y1, x0:x1],
                    output_shape=(224, 224))
                part_imgs.append(x_res)
        part_imgs = F.concat(tuple(part_imgs), axis=0)
        part_features = self.backbone_tail(self.backbone(part_imgs))

        part_feature = part_features.reshape((batch, self.top_n, -1))
        part_feature = part_feature[:, :self.num_cat, :]
        part_feature = part_feature.reshape((batch, -1))

        raw_features = self.backbone_tail(raw_pre_features)

        concat_out = F.concat((part_feature, raw_features), axis=1)
        concat_logits = self.concat_net(concat_out)

        if self.aux:
            raw_logits = self.backbone_classifier(raw_features)
            part_logits = self.partcls_net(part_features).reshape((batch, self.top_n, -1))
            return concat_logits, raw_logits, part_logits, top_n_prob
        else:
            return concat_logits 

# 需要导入模块: import numpy [as 别名]
# 或者: from numpy import take_along_axis [as 别名]
def testTakeAlongAxisIssue1521(self):
    # https://github.com/google/jax/issues/1521
    idx = lnp.repeat(lnp.arange(3), 10).reshape((30, 1))

    def f(x):
      y = x * lnp.arange(3.).reshape((1, 3))
      return lnp.take_along_axis(y, idx, -1).sum()

    check_grads(f, (1.,), order=1)