def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
scores_inds = flow_exp.argsort(scores, dim=0, descending=True)
boxes = flow._C.gather(boxes, scores_inds, axis=0)
_nms_op = (flow_exp.builtin_op("nms").Input("in").Output("out").Attr(
"iou_threshold", iou_threshold).Attr("keep_n", -1).Build())
keep = _nms_op(boxes)[0]
index = flow_exp.squeeze(flow_exp.argwhere(keep), dim=[1])
return flow._C.gather(scores_inds, index, axis=0)
def _test_argwhere_with_random_data(test_case, ndim, placement, sbp):
dims = [random(1, 3) * 8 for _ in range(ndim)]
x = random_tensor(ndim, *dims).to_global(placement=placement, sbp=sbp)
# PyTorch has no argwhere before v1.11, so we use nonzero instead of argwhere for PyTorch
# y = torch.argwhere(x)
y = x.clone()
y.oneflow = flow.argwhere(x.oneflow)
y.pytorch = torch_ori.nonzero(x.pytorch)
return y
def _test_argwhere(test_case, shape, device):
np_input = np.random.randn(*shape)
input = flow.tensor(np_input,
dtype=flow.float32,
device=flow.device(device))
of_out = flow.argwhere(input)
np_out = np.argwhere(np_input)
test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 0.0001, 0.0001))
test_case.assertTrue(np.array_equal(of_out.numpy().shape, np_out.shape))
def masked_select_op(input, mask):
"""
Returns a new 1-D tensor which indexes the input tensor according to the boolean mask mask which is a BoolTensor(In oneFlow BoolTensor is replaced by Int8Tensor).
The shapes of the mask tensor and the input tensor don’t need to match, but they must be broadcastable.
Args:
input (Tensor): the input tensor.
mask (Tensor): the tensor containing the binary mask to index with
For example:
.. code-block:: python
>>> import oneflow as flow
>>> import numpy as np
>>> input = flow.tensor(np.array([[-0.4620, 0.3139], [0.3898, -0.7197], [0.0478, -0.1657]]), dtype=flow.float32)
>>> mask = input.gt(0.05)
>>> out = flow.masked_select(input, mask)
>>> out
tensor([0.3139, 0.3898], dtype=oneflow.float32)
"""
assert len(input.shape) == len(
mask.shape
), f"The dim of masked_select module's inputs can not match, please check!"
broadcast_like_shape = []
broadcast_x_axes = []
broadcast_mask_axes = []
for i in range(len(input.shape)):
max_dim = max(input.shape[i], mask.shape[i])
broadcast_like_shape.append(max_dim)
if max_dim != input.shape[i]:
broadcast_x_axes.append(i)
if max_dim != mask.shape[i]:
broadcast_mask_axes.append(i)
broadcast_like_tensor = flow.zeros(tuple(broadcast_like_shape),
dtype=flow.float32,
device=input.device)
broadcast_like_tensor.requires_grad = input.requires_grad or mask.requires_grad
if len(broadcast_x_axes) != 0:
input = flow.broadcast_like(input,
broadcast_like_tensor,
broadcast_axes=tuple(broadcast_x_axes))
if len(broadcast_mask_axes) != 0:
mask = flow.broadcast_like(mask,
broadcast_like_tensor,
broadcast_axes=tuple(broadcast_mask_axes))
mask = mask.to(dtype=input.dtype)
res = flow._C.mul(input, mask)
indices = flow.argwhere(res)
gather_res = flow._C.gather_nd(res, indices)
return gather_res.flatten()
def _argwhere(self):
return flow.argwhere(self)
def nms_op(boxes, scores, iou_threshold: float):
score_inds = flow.argsort(scores, dim=0, descending=True)
boxes = flow._C.gather(boxes, score_inds, axis=0)
keep = flow._C.nms(boxes, iou_threshold)
index = flow.squeeze(flow.argwhere(keep), dim=[1])
return flow._C.gather(score_inds, index, axis=0)
def argwhere_fn(x: flow.typing.Numpy.Placeholder(
x.shape, dtype=data_type)) -> flow.typing.ListNumpy:
return flow.argwhere(x, dtype=out_data_type)
def do_argwhere(x_blob):
with flow.scope.placement(device_type, "0:0"):
return flow.argwhere(x_blob, dtype=out_data_type)
