|
import torch |
|
import numpy as np |
|
from typing import Tuple, Optional |
|
|
|
|
|
def find_runs(x: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
|
"""Find runs of consecutive items in an array.""" |
|
|
|
|
|
x = np.asanyarray(x) |
|
if x.ndim != 1: |
|
raise ValueError("only 1D array supported") |
|
n = x.shape[0] |
|
|
|
|
|
if n == 0: |
|
return np.array([]), np.array([]), np.array([]) |
|
else: |
|
|
|
loc_run_start = np.empty(n, dtype=bool) |
|
loc_run_start[0] = True |
|
np.not_equal(x[:-1], x[1:], out=loc_run_start[1:]) |
|
run_starts = np.nonzero(loc_run_start)[0] |
|
|
|
|
|
run_values = x[loc_run_start] |
|
|
|
|
|
run_lengths = np.diff(np.append(run_starts, n)) |
|
|
|
return run_values, run_starts, run_lengths |
|
|
|
|
|
def compute_mask_indices( |
|
shape: Tuple[int, int], |
|
padding_mask: Optional[torch.Tensor], |
|
mask_prob: float, |
|
mask_length: int, |
|
mask_type: str = "static", |
|
mask_other: float = 0.0, |
|
min_masks: int = 0, |
|
no_overlap: bool = False, |
|
min_space: int = 0, |
|
) -> np.ndarray: |
|
""" |
|
Computes random mask spans for a given shape |
|
Args: |
|
shape: the the shape for which to compute masks. |
|
should be of size 2 where first element is batch size and 2nd is timesteps |
|
padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements |
|
mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by |
|
number of timesteps divided by length of mask span to mask approximately this percentage of all elements. |
|
however due to overlaps, the actual number will be smaller (unless no_overlap is True) |
|
mask_type: how to compute mask lengths |
|
static = fixed size |
|
uniform = sample from uniform distribution [mask_other, mask_length*2] |
|
normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element |
|
poisson = sample from possion distribution with lambda = mask length |
|
min_masks: minimum number of masked spans |
|
no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping |
|
min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans |
|
""" |
|
bsz, all_sz = shape |
|
mask = np.full((bsz, all_sz), False) |
|
|
|
all_num_mask = int( |
|
|
|
mask_prob * all_sz / float(mask_length) |
|
+ np.random.rand() |
|
) |
|
|
|
all_num_mask = max(min_masks, all_num_mask) |
|
|
|
mask_idcs = [] |
|
for i in range(bsz): |
|
if padding_mask is not None: |
|
sz = all_sz - padding_mask[i].long().sum().item() |
|
num_mask = int( |
|
|
|
mask_prob * sz / float(mask_length) |
|
+ np.random.rand() |
|
) |
|
num_mask = max(min_masks, num_mask) |
|
else: |
|
sz = all_sz |
|
num_mask = all_num_mask |
|
|
|
if mask_type == "static": |
|
lengths = np.full(num_mask, mask_length) |
|
elif mask_type == "uniform": |
|
lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask) |
|
elif mask_type == "normal": |
|
lengths = np.random.normal(mask_length, mask_other, size=num_mask) |
|
lengths = [max(1, int(round(x))) for x in lengths] |
|
elif mask_type == "poisson": |
|
lengths = np.random.poisson(mask_length, size=num_mask) |
|
lengths = [int(round(x)) for x in lengths] |
|
else: |
|
raise Exception("unknown mask selection " + mask_type) |
|
|
|
if sum(lengths) == 0: |
|
lengths[0] = min(mask_length, sz - 1) |
|
|
|
if no_overlap: |
|
mask_idc = [] |
|
|
|
def arrange(s, e, length, keep_length): |
|
span_start = np.random.randint(s, e - length) |
|
mask_idc.extend(span_start + i for i in range(length)) |
|
|
|
new_parts = [] |
|
if span_start - s - min_space >= keep_length: |
|
new_parts.append((s, span_start - min_space + 1)) |
|
if e - span_start - keep_length - min_space > keep_length: |
|
new_parts.append((span_start + length + min_space, e)) |
|
return new_parts |
|
|
|
parts = [(0, sz)] |
|
min_length = min(lengths) |
|
for length in sorted(lengths, reverse=True): |
|
lens = np.fromiter( |
|
(e - s if e - s >= length + min_space else 0 for s, e in parts), |
|
np.int, |
|
) |
|
l_sum = np.sum(lens) |
|
if l_sum == 0: |
|
break |
|
probs = lens / np.sum(lens) |
|
c = np.random.choice(len(parts), p=probs) |
|
s, e = parts.pop(c) |
|
parts.extend(arrange(s, e, length, min_length)) |
|
mask_idc = np.asarray(mask_idc) |
|
else: |
|
min_len = min(lengths) |
|
if sz - min_len <= num_mask: |
|
min_len = sz - num_mask - 1 |
|
|
|
mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) |
|
|
|
mask_idc = np.asarray( |
|
[ |
|
mask_idc[j] + offset |
|
for j in range(len(mask_idc)) |
|
for offset in range(lengths[j]) |
|
] |
|
) |
|
|
|
mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) |
|
|
|
min_len = min([len(m) for m in mask_idcs]) |
|
batch_indexes, starts, ends = [], [], [] |
|
for i, mask_idc in enumerate(mask_idcs): |
|
if len(mask_idc) > min_len: |
|
mask_idc = np.random.choice(mask_idc, min_len, replace=False) |
|
mask[i, mask_idc] = True |
|
vals, run_starts, run_lengths = find_runs(mask[i]) |
|
start_indices, lengths = run_starts[vals == True], run_lengths[vals == True] |
|
starts.append(start_indices) |
|
ends.append(start_indices + lengths) |
|
batch_indexes.append(np.zeros([len(start_indices)]) + i) |
|
return ( |
|
mask, |
|
np.concatenate(starts).astype(np.int64), |
|
np.concatenate(ends).astype(np.int64), |
|
np.concatenate(batch_indexes).astype(np.int64), |
|
) |
|
|
