improve documentation for sampler
Summary:
(fix https://github.com/facebookresearch/detectron2/issues/431)
Pull Request resolved:
@@ -6,8 +6,11 @@ __all__ = ["subsample_labels"]
def subsample_labels(labels, num_samples, positive_fraction, bg_label):
"""
- Return `num_samples` random samples from `labels`, with a fraction of
- positives no larger than `positive_fraction`.
+ 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:
@@ -26,8 +29,7 @@ def subsample_labels(labels, num_samples, positive_fraction, bg_label):
Returns:
pos_idx, neg_idx (Tensor):
- 1D indices. The total number of indices is `num_samples` if possible.
- The fraction of positive indices is `positive_fraction` if possible.
+ 1D vector of indices. The total length of both is `num_samples` or fewer.
"""
positive = torch.nonzero((labels != -1) & (labels != bg_label)).squeeze(1)
negative = torch.nonzero(labels == bg_label).squeeze(1)
