def prune(scores: mx.nd.NDArray, finished: mx.nd.NDArray,
inf_array: mx.nd.NDArray, beam_size: int,
prune_threshold: float) -> mx.nd.NDArray:
"""
Returns a 0-1 array indicating which hypotheses are inactive based on pruning.
Finished hypotheses that have a score worse than prune_threshold from the best scoring hypotheses
are marked as inactive.
:param scores: Hypotheses scores. Shape: (batch * beam, 1).
:param finished: 0-1 array indicating which hypotheses are finished. Shape: (batch * beam,).
:param inf_array: Auxiliary array filled with infinity. Shape: (batch * beam,).
:param beam_size: Beam size.
:param prune_threshold: Pruning threshold.
:return NDArray of inactive items. Shape(batch * beam,).
"""
scores = scores.reshape((-1, beam_size))
finished = finished.reshape((-1, beam_size))
inf_array = inf_array.reshape((-1, beam_size))
# best finished scores. Shape: (batch, 1)
best_finished_scores = mx.nd.where(finished, scores,
inf_array).min(axis=1, keepdims=True)
inactive = mx.nd.cast((scores - best_finished_scores) > prune_threshold,
dtype='int32').reshape((-1))
return inactive
def bbox_overlaps(anchors: mx.nd.NDArray, gt: mx.nd.NDArray):
"""
Get IoU of the anchors and ground truth bounding boxes.
The shape of anchors and gt should be (N, 4) and (M, 4)
So the shape of return value is (N, M)
"""
N, M = anchors.shape[0], gt.shape[0]
anchors_mat = anchors.reshape((N, 1, 4)).broadcast_to((N, M, 4)).reshape(
(-1, 4))
gt_mat = gt.reshape((1, M, 4)).broadcast_to((N, M, 4)).reshape((-1, 4))
# inter
x0 = nd.max(nd.stack(anchors_mat[:, 0], gt_mat[:, 0]), axis=0)
y0 = nd.max(nd.stack(anchors_mat[:, 1], gt_mat[:, 1]), axis=0)
x1 = nd.min(nd.stack(anchors_mat[:, 2], gt_mat[:, 2]), axis=0)
y1 = nd.min(nd.stack(anchors_mat[:, 3], gt_mat[:, 3]), axis=0)
inter = _get_area(
nd.concatenate([
x0.reshape((-1, 1)),
y0.reshape((-1, 1)),
x1.reshape((-1, 1)),
y1.reshape((-1, 1))
],
axis=1))
outer = _get_area(anchors_mat) + _get_area(gt_mat) - inter
iou = inter / outer
iou = iou.reshape((N, M))
return iou
def topk(scores: mx.nd.NDArray,
offset: mx.nd.NDArray,
k: int) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
"""
Get the lowest k elements per sentence from a `scores` matrix.
At the first timestep, the shape of scores is (batch, target_vocabulary_size).
At subsequent steps, the shape is (batch * k, target_vocabulary_size).
:param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
:param offset: Array (shape: batch_size * k) containing offsets to add to the hypothesis indices in batch decoding.
:param k: The number of smallest scores to return.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
# Compute the batch size from the offsets and k. We don't know the batch size because it is
# either 1 (at timestep 1) or k (at timesteps 2+).
# (batch_size, beam_size * target_vocab_size)
batch_size = int(offset.shape[-1] / k)
folded_scores = scores.reshape((batch_size, -1))
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_scores, axis=1, k=k, ret_typ='both', is_ascend=True)
indices = mx.nd.cast(indices, 'int32').reshape((-1,))
best_hyp_indices, best_word_indices = mx.nd.unravel_index(indices, shape=(batch_size * k, scores.shape[-1]))
if batch_size > 1:
# Offsetting the indices to match the shape of the scores matrix
best_hyp_indices += offset
values = values.reshape((-1, 1))
return best_hyp_indices, best_word_indices, values
def topk(
scores: mx.nd.NDArray, k: int, offset: mx.nd.NDArray
) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
"""
Get the lowest k elements per sentence from a `scores` matrix.
:param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
:param k: The number of smallest scores to return.
:param offset: Array to add to the hypothesis indices for offsetting in batch decoding.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
# (batch_size, beam_size * target_vocab_size)
folded_scores = scores.reshape((-1, k * scores.shape[-1]))
batch_size = folded_scores.shape[0]
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_scores,
axis=1,
k=k,
ret_typ='both',
is_ascend=True)
indices = mx.nd.cast(indices, 'int32').reshape((-1, ))
best_hyp_indices, best_word_indices = mx.nd.unravel_index(
indices, scores.shape)
if batch_size > 1:
# Offsetting the indices to match the shape of the scores matrix
best_hyp_indices += offset
values = values.reshape((-1, 1))
return best_hyp_indices, best_word_indices, values
def smallest_k_mx(
matrix: mx.nd.NDArray,
k: int,
only_first_row: bool = False
) -> Tuple[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
Find the smallest elements in a NDarray.
:param matrix: Any matrix.
:param k: The number of smallest elements to return.
:param only_first_row: If True the search is constrained to the first row of the matrix.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
if only_first_row:
folded_matrix = mx.nd.reshape(matrix[0], shape=(-1))
else:
folded_matrix = matrix.reshape(-1)
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_matrix,
axis=None,
k=k,
ret_typ='both',
is_ascend=True,
dtype="int32")
indices = np.unravel_index(
indices.reshape(-1).astype(np.int64).asnumpy(), matrix.shape)
logger.info(indices)
return indices, values
def update(self, preds: mx.nd.NDArray, labels: mx.nd.NDArray):
"""
Update confusion table used to compute scores.
:param preds: predicted classes for samples in batch.
:type preds: mx.ndarray.NDArray
:param labels: actual labels for samples in batch.
:type labels: mx.ndarray.NDArray
"""
preds = preds.reshape((-1, 1))
labels = labels.reshape((-1, 1))
preds_binary = (preds == self.all_labels)
labels_binary = (labels == self.all_labels)
self.confusion_matrix += mx.nd.dot(labels_binary.T, preds_binary)
def topk(
scores: mx.nd.NDArray, k: int, batch_size: int, offset: mx.nd.NDArray,
use_mxnet_topk: bool
) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
"""
Get the lowest k elements per sentence from a `scores` matrix.
:param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
:param k: The number of smallest scores to return.
:param batch_size: Number of sentences being decoded at once.
:param offset: Array to add to the hypothesis indices for offsetting in batch decoding.
:param use_mxnet_topk: True to use the mxnet implementation or False to use the numpy one.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
# (batch_size, beam_size * target_vocab_size)
folded_scores = scores.reshape((batch_size, k * scores.shape[-1]))
if use_mxnet_topk:
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_scores,
axis=1,
k=k,
ret_typ='both',
is_ascend=True)
best_hyp_indices, best_word_indices = mx.nd.array(np.unravel_index(
indices.astype(np.int32).asnumpy().ravel(), scores.shape),
dtype='int32',
ctx=scores.context)
else:
folded_scores = folded_scores.asnumpy()
# Get the scores
# Indexes into folded_scores: (batch_size, beam_size)
flat_idxs = np.argpartition(folded_scores, range(k))[:, :k]
# Score values: (batch_size, beam_size)
values = mx.nd.array(
folded_scores[np.arange(folded_scores.shape[0])[:, None],
flat_idxs],
ctx=scores.context)
best_hyp_indices, best_word_indices = mx.nd.array(np.unravel_index(
flat_idxs.ravel(), scores.shape),
dtype='int32',
ctx=scores.context)
if batch_size > 1:
# Offsetting the indices to match the shape of the scores matrix
best_hyp_indices += offset
values = values.reshape((-1, 1))
return best_hyp_indices, best_word_indices, values
def smallest_k_mx_batched(
matrix: mx.nd.NDArray,
k: int,
batch_size: int,
offset: mx.nd.NDArray,
only_first_row: bool = False
) -> Tuple[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
Find the smallest elements in a NDarray.
:param matrix: Any matrix.
:param k: The number of smallest elements to return.
:param only_first_row: If True the search is constrained to the first row of the matrix.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
if only_first_row:
folded_matrix = matrix.reshape((-4, batch_size, -1, 0, 0))
folded_matrix = folded_matrix[:, 0, :, :]
else:
folded_matrix = matrix
folded_matrix = folded_matrix.reshape((batch_size, -1))
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_matrix,
axis=1,
k=k,
ret_typ='both',
is_ascend=True,
dtype="int32")
# best_hyp_indices, best_hyp_pos_indices, best_word_indices = mx.nd.array(np.unravel_index(indices.astype(np.int32).asnumpy().ravel(), matrix.shape),
# dtype='int32',
# ctx=matrix.context)
best_hyp_indices, best_hyp_pos_indices, best_word_indices = mx.nd.array(
mx.nd.unravel_index(
indices.astype(np.int32).reshape(-1), matrix.shape),
dtype='int32',
ctx=matrix.context)
if batch_size > 1:
best_hyp_indices += offset
# assert(mx.nd.nansum(values.reshape(-1) - matrix[best_hyp_indices, best_hyp_pos_indices, best_word_indices]) == 0)
return (best_hyp_indices, best_hyp_pos_indices,
best_word_indices), values.reshape(-1)
def topk(scores: mx.nd.NDArray,
k: int,
batch_size: int,
offset: mx.nd.NDArray,
use_mxnet_topk: bool) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
"""
Get the lowest k elements per sentence from a `scores` matrix.
:param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
:param k: The number of smallest scores to return.
:param batch_size: Number of sentences being decoded at once.
:param offset: Array to add to the hypothesis indices for offsetting in batch decoding.
:param use_mxnet_topk: True to use the mxnet implementation or False to use the numpy one.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
# (batch_size, beam_size * target_vocab_size)
folded_scores = scores.reshape((batch_size, k * scores.shape[-1]))
if use_mxnet_topk:
# pylint: disable=unbalanced-tuple-unpacking
values, indices = mx.nd.topk(folded_scores, axis=1, k=k, ret_typ='both', is_ascend=True)
best_hyp_indices, best_word_indices = mx.nd.array(np.unravel_index(indices.astype(np.int32).asnumpy().ravel(),
scores.shape),
dtype='int32',
ctx=scores.context)
else:
folded_scores = folded_scores.asnumpy()
# Get the scores
# Indexes into folded_scores: (batch_size, beam_size)
flat_idxs = np.argpartition(folded_scores, range(k))[:, :k]
# Score values: (batch_size, beam_size)
values = mx.nd.array(folded_scores[np.arange(folded_scores.shape[0])[:, None], flat_idxs], ctx=scores.context)
best_hyp_indices, best_word_indices = mx.nd.array(np.unravel_index(flat_idxs.ravel(), scores.shape),
dtype='int32', ctx=scores.context)
if batch_size > 1:
# Offsetting the indices to match the shape of the scores matrix
best_hyp_indices += offset
values = values.reshape((-1, 1))
return best_hyp_indices, best_word_indices, values
def numpy_topk(
scores: mx.nd.NDArray, k: int, offset: mx.nd.NDArray
) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
"""
Get the lowest k elements per sentence from a `scores` matrix using an intermediary Numpy conversion.
This should be equivalent to sockeye.utils.topk() and is used as a comparative implementation in testing.
:param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
:param k: The number of smallest scores to return.
:param offset: Array to add to the hypothesis indices for offsetting in batch decoding.
:return: The row indices, column indices and values of the k smallest items in matrix.
"""
# (batch_size, beam_size * target_vocab_size)
folded_scores = scores.reshape((-1, k * scores.shape[-1]))
batch_size = folded_scores.shape[0]
folded_scores = folded_scores.asnumpy()
# Get the scores
# Indexes into folded_scores: (batch_size, beam_size)
flat_idxs = np.argpartition(folded_scores, range(k))[:, :k]
# Score values: (batch_size, beam_size)
values = mx.nd.array(folded_scores[np.arange(folded_scores.shape[0])[:,
None],
flat_idxs],
ctx=scores.context)
best_hyp_indices, best_word_indices = mx.nd.array(np.unravel_index(
flat_idxs.ravel(), scores.shape),
dtype='int32',
ctx=scores.context)
if batch_size > 1:
# Offsetting the indices to match the shape of the scores matrix
best_hyp_indices += offset
values = values.reshape((-1, 1))
return best_hyp_indices, best_word_indices, values
