def _cumsum(x, axis=0, exclusive=0, reverse=0, name=''):
if isinstance(axis, int):
axis = bF.to_tensor(axis)
return bF.cumsum(x,
axis,
exclusive=exclusive,
reverse=reverse,
debugContext=name)
def _top_k(input, k=1, sorted=True, name=''):
dim = input.shape.ndims - 1
assert input.pureShape[-1] >= k
k = bF.to_tensor(k, dtype='INT64')
values, order = bF.topk(input,
k,
sorted=sorted,
dim=dim,
debugContext=name)
return values, order
def nms(input_scores,
input_boxes,
threshold=0.7,
numDetections=300,
score_threshold=None,
debugContext=''):
load_lib()
input_scores = input_scores.cast('FLOAT')
input_boxes = input_boxes.cast('FLOAT')
valid_area_mask = bF.transpose(get_valid_area_mask(input_boxes),
[1, 0]) # 1,n
input_scores = input_scores + 1e-6 # if score==0, proposals will be ignored
local_input_scores = bF.identity(input_scores * valid_area_mask,
debugContext=debugContext).detach()
local_input_boxes = bF.identity(input_boxes,
debugContext=debugContext).detach()
if local_input_scores.shape.ndims == 1:
local_input_scores = local_input_scores.unsqueeze(0)
if local_input_boxes.shape.ndims == 2:
local_input_boxes = local_input_boxes.unsqueeze(0)
assert local_input_boxes.pureShape[0] == 1, 'only implemented batch=1'
if score_threshold is not None:
assert isinstance(score_threshold, float)
local_mask = bF.greater(
local_input_scores,
bF.to_tensor(score_threshold, dtype=local_input_scores.dtype))
local_mask = bF.cast(local_mask, target_type=local_input_scores.dtype)
local_input_scores = local_input_scores * local_mask
with bF.name_scope("nms"):
out = bF.get_builder().customOp(opName="nms",
opVersion=1,
domain="ai.graphcore",
inputs=[
local_input_scores.getIpuIndex(),
local_input_boxes.getIpuIndex()
],
attributes={
"threshold": threshold,
"numDetections": numDetections
},
numOutputs=3,
name="nmsCustomOp")
#
_, output_boxes, output_keep = out[0], bF.TTensor(out[1]), bF.TTensor(
out[2])
targetType = input_scores.dtype
roiKeeps_flag = bF.cast(bF.greater(
output_keep, bF.constant(np.asarray(-1, dtype=np.int32))),
target_type='INT32')
num_valids = bF.reduceSum(roiKeeps_flag, axes=[1])
roiKeeps_flag = bF.cast(roiKeeps_flag, target_type=targetType)
roiKeeps_flag = bF.unsqueeze(roiKeeps_flag, [-1])
output_boxes = bF.mul([output_boxes, roiKeeps_flag])
return output_boxes, output_keep, num_valids
def slice(input_, begin, size, name=''):
axes = list(builtins.range(len(begin)))
input_shape = input_.pureShape
assert len(begin) == len(size)
ends = []
for begin_ele, size_ele, shape_ele in zip(begin, size, input_shape):
if size_ele == -1:
ends.append(shape_ele)
else:
ends.append(begin_ele + size_ele)
if isinstance(begin, list):
begin = bF.to_tensor(begin)
if isinstance(ends, list):
ends = bF.to_tensor(ends)
if isinstance(axes, list):
axes = bF.to_tensor(axes)
return bF.gc_slice(input_,
axes,
starts=begin,
ends=ends,
debugContext=name)
def minimum(x, y, name=""):
x = bF.to_tensor(x)
y = bF.to_tensor(y)
return bF.min([x, y], debugContext=name)
def less(x, y, name=''):
x = bF.to_tensor(x)
y = bF.to_tensor(y)
return bF.less(x, y, debugContext=name)
def greater(x, y, name=''):
x = bF.to_tensor(x)
y = bF.to_tensor(y)
return bF.greater(x, y, debugContext=name)