Fastai v2 audio

fastai fastai dev
26

I’ve managed to get SpecAugment for batches working with two different methods:

Method 1

Uses a for loop but is decorated as a torch script to speed it up

@torch.jit.script
def spec_aug_loop(batch:Tensor, size:int=20):
    bsg = batch.clone()
    max_y = bsg.shape[-2]-size-1
    for i in range(bsg.shape[0]):
        s = bsg[i, :]
        m = s.flatten(-2).mean()
        r = torch.randint(0,max_y,(1,)).squeeze().int()
        s[:, r:r+size, :] = m
    return bsg

2.83 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Method 2

Using masks to replace the values in random places

@torch.jit.script
def spec_augment(batch:Tensor, size:int=20):
    bsg = batch.clone()
    bs, _, x, y = bsg.shape
    max_y = y-size-torch.tensor(1)
    m = torch.arange(y).repeat(x*bs).view(bs,-1)
    rs = torch.randint(0,max_y,(1,bs)).squeeze()[None].t()
    gpumask = ((m > rs)) & (m < (rs+size))
    gpumask = gpumask.view(bs,x,-1)[:,None,:]
    bsg[gpumask] = torch.tensor(0)
    return bsg

7.26 ms ± 81.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

The clear winner is method 1 but I believe that there must be something I can do to speed up method 2

3 Likes