import torch
if torch.__version__ >= '1.8':
import torch_npu
import numpy as np
from detectron2.layers import nonzero_tuple
__all__ = ["subsample_labels", "subsample_labels_all"]
def subsample_labels_bk(
labels: torch.Tensor, num_samples: int, positive_fraction: float, bg_label: int
):
"""
Return `num_samples` (or fewer, if not enough found)
random samples from `labels` which is a mixture of positives & negatives.
It will try to return as many positives as possible without
exceeding `positive_fraction * num_samples`, and then try to
fill the remaining slots with negatives.
Args:
labels (Tensor): (N, ) label vector with values:
* -1: ignore
* bg_label: background ("negative") class
* otherwise: one or more foreground ("positive") classes
num_samples (int): The total number of labels with value >= 0 to return.
Values that are not sampled will be filled with -1 (ignore).
positive_fraction (float): The number of subsampled labels with values > 0
is `min(num_positives, int(positive_fraction * num_samples))`. The number
of negatives sampled is `min(num_negatives, num_samples - num_positives_sampled)`.
In order words, if there are not enough positives, the sample is filled with
negatives. If there are also not enough negatives, then as many elements are
sampled as is possible.
bg_label (int): label index of background ("negative") class.
Returns:
pos_idx, neg_idx (Tensor):
1D vector of indices. The total length of both is `num_samples` or fewer.
"""
positive = nonzero_tuple((labels != -1) & (labels != bg_label))[0]
negative = nonzero_tuple(labels == bg_label)[0]
num_pos = int(num_samples * positive_fraction)
num_pos = min(positive.numel(), num_pos)
num_neg = num_samples - num_pos
num_neg = min(negative.numel(), num_neg)
perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]
pos_idx = positive[perm1]
neg_idx = negative[perm2]
return pos_idx, neg_idx
def subsample_labels(
labels: torch.Tensor, num_samples: int, positive_fraction: float, bg_label: int
):
"""
Return `num_samples` (or fewer, if not enough found)
random samples from `labels` which is a mixture of positives & negatives.
It will try to return as many positives as possible without
exceeding `positive_fraction * num_samples`, and then try to
fill the remaining slots with negatives.
Args:
labels (Tensor): (N, ) label vector with values:
* -1: ignore
* bg_label: background ("negative") class
* otherwise: one or more foreground ("positive") classes
num_samples (int): The total number of labels with value >= 0 to return.
Values that are not sampled will be filled with -1 (ignore).
positive_fraction (float): The number of subsampled labels with values > 0
is `min(num_positives, int(positive_fraction * num_samples))`. The number
of negatives sampled is `min(num_negatives, num_samples - num_positives_sampled)`.
In order words, if there are not enough positives, the sample is filled with
negatives. If there are also not enough negatives, then as many elements are
sampled as is possible.
bg_label (int): label index of background ("negative") class.
Returns:
pos_idx, neg_idx (Tensor):
1D vector of indices. The total length of both is `num_samples` or fewer.
"""
positive = (labels != -1) & (labels != bg_label)
negative = labels == bg_label
num_pos = int(num_samples * positive_fraction)
pos_idx, pos_mask = torch_npu.npu_random_choice_with_mask(positive, num_pos)
num_pos = min(pos_mask.sum(), num_pos)
neg_idx, neg_mask = torch_npu.npu_random_choice_with_mask(negative, num_samples)
num_neg = min(neg_mask.sum(), num_samples - num_pos)
return pos_idx.reshape((-1,))[:num_pos], neg_idx.reshape((-1,))[:num_neg]
def subsample_labels_all(
labels: torch.Tensor, num_samples: int, positive_fraction: float, bg_label: int
):
"""
Return `num_samples` (or fewer, if not enough found)
random samples from `labels` which is a mixture of positives & negatives.
It will try to return as many positives as possible without
exceeding `positive_fraction * num_samples`, and then try to
fill the remaining slots with negatives.
Args:
labels (Tensor): (N, ) label vector with values:
* -1: ignore
* bg_label: background ("negative") class
* otherwise: one or more foreground ("positive") classes
num_samples (int): The total number of labels with value >= 0 to return.
Values that are not sampled will be filled with -1 (ignore).
positive_fraction (float): The number of subsampled labels with values > 0
is `min(num_positives, int(positive_fraction * num_samples))`. The number
of negatives sampled is `min(num_negatives, num_samples - num_positives_sampled)`.
In order words, if there are not enough positives, the sample is filled with
negatives. If there are also not enough negatives, then as many elements are
sampled as is possible.
bg_label (int): label index of background ("negative") class.
Returns:
pos_idx, neg_idx (Tensor):
1D vector of indices. The total length of both is `num_samples` or fewer.
"""
positive_mask = ((labels != -1) & (labels != bg_label)).int()
negative_mask = (labels == bg_label).int()
positive = nonzero_tuple(positive_mask)[0]
negative = nonzero_tuple(negative_mask)[0]
num_pos = int(num_samples * positive_fraction)
num_pos = min(positive.numel(), num_pos)
num_neg = num_samples - num_pos
num_neg = min(negative.numel(), num_neg)
num_pos_unuse = positive.numel() - num_pos
num_neg_unuse = negative.numel() - num_neg
if num_pos_unuse != 0:
perm1_un = torch.randperm(
positive.numel(), device=positive.device
)[:num_pos_unuse]
positive_unuse = positive[perm1_un]
positive_mask[positive_unuse.long()] = 0
if num_neg_unuse != 0:
perm2_un = torch.randperm(
negative.numel(), device=negative.device
)[:num_neg_unuse]
negative_unuse = negative[perm2_un]
negative_mask[negative_unuse.long()] = 0
sampled_idxs_mask = positive_mask + negative_mask
sampled_idxs = (nonzero_tuple(sampled_idxs_mask)[0]).long()
return sampled_idxs