Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # -*- coding: utf-8 -*- | |
| # @Author : Xinhao Mei @CVSSP, University of Surrey | |
| # @E-mail : x.mei@surrey.ac.uk | |
| """ | |
| Implemenation of SpecAugment++, | |
| Adapated from Qiuqiang Kong's trochlibrosa: | |
| https://github.com/qiuqiangkong/torchlibrosa/blob/master/torchlibrosa/augmentation.py | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class DropStripes: | |
| def __init__(self, dim, drop_width, stripes_num): | |
| """ Drop stripes. | |
| args: | |
| dim: int, dimension along which to drop | |
| drop_width: int, maximum width of stripes to drop | |
| stripes_num: int, how many stripes to drop | |
| """ | |
| super(DropStripes, self).__init__() | |
| assert dim in [2, 3] # dim 2: time; dim 3: frequency | |
| self.dim = dim | |
| self.drop_width = drop_width | |
| self.stripes_num = stripes_num | |
| def __call__(self, input): | |
| """input: (batch_size, channels, time_steps, freq_bins)""" | |
| assert input.ndimension() == 4 | |
| batch_size = input.shape[0] | |
| total_width = input.shape[self.dim] | |
| for n in range(batch_size): | |
| self.transform_slice(input[n], total_width) | |
| return input | |
| def transform_slice(self, e, total_width): | |
| """ e: (channels, time_steps, freq_bins)""" | |
| for _ in range(self.stripes_num): | |
| distance = torch.randint(low=0, high=self.drop_width, size=(1,))[0] | |
| bgn = torch.randint(low=0, high=total_width - distance, size=(1,))[0] | |
| if self.dim == 2: | |
| e[:, bgn: bgn + distance, :] = 0 | |
| elif self.dim == 3: | |
| e[:, :, bgn: bgn + distance] = 0 | |
| class MixStripes: | |
| def __init__(self, dim, mix_width, stripes_num): | |
| """ Mix stripes | |
| args: | |
| dim: int, dimension along which to mix | |
| mix_width: int, maximum width of stripes to mix | |
| stripes_num: int, how many stripes to mix | |
| """ | |
| super(MixStripes, self).__init__() | |
| assert dim in [2, 3] | |
| self.dim = dim | |
| self.mix_width = mix_width | |
| self.stripes_num = stripes_num | |
| def __call__(self, input): | |
| """input: (batch_size, channel, time_steps, freq_bins)""" | |
| assert input.ndimension() == 4 | |
| batch_size = input.shape[0] | |
| total_width = input.shape[self.dim] | |
| rand_sample = input[torch.randperm(batch_size)] | |
| for i in range(batch_size): | |
| self.transform_slice(input[i], rand_sample[i], total_width) | |
| return input | |
| def transform_slice(self, input, random_sample, total_width): | |
| for _ in range(self.stripes_num): | |
| distance = torch.randint(low=0, high=self.mix_width, size=(1,))[0] | |
| bgn = torch.randint(low=0, high=total_width - distance, size=(1,))[0] | |
| if self.dim == 2: | |
| input[:, bgn: bgn + distance, :] = 0.5 * input[:, bgn: bgn + distance, :] + \ | |
| 0.5 * random_sample[:, bgn: bgn + distance, :] | |
| elif self.dim == 3: | |
| input[:, :, bgn: bgn + distance] = 0.5 * input[:, :, bgn: bgn + distance] + \ | |
| 0.5 * random_sample[:, :, bgn: bgn + distance] | |
| class CutStripes: | |
| def __init__(self, dim, cut_width, stripes_num): | |
| """ Cutting stripes with another randomly selected sample in mini-batch. | |
| args: | |
| dim: int, dimension along which to cut | |
| cut_width: int, maximum width of stripes to cut | |
| stripes_num: int, how many stripes to cut | |
| """ | |
| super(CutStripes, self).__init__() | |
| assert dim in [2, 3] | |
| self.dim = dim | |
| self.cut_width = cut_width | |
| self.stripes_num = stripes_num | |
| def __call__(self, input): | |
| """input: (batch_size, channel, time_steps, freq_bins)""" | |
| assert input.ndimension() == 4 | |
| batch_size = input.shape[0] | |
| total_width = input.shape[self.dim] | |
| rand_sample = input[torch.randperm(batch_size)] | |
| for i in range(batch_size): | |
| self.transform_slice(input[i], rand_sample[i], total_width) | |
| return input | |
| def transform_slice(self, input, random_sample, total_width): | |
| for _ in range(self.stripes_num): | |
| distance = torch.randint(low=0, high=self.cut_width, size=(1,))[0] | |
| bgn = torch.randint(low=0, high=total_width - distance, size=(1,))[0] | |
| if self.dim == 2: | |
| input[:, bgn: bgn + distance, :] = random_sample[:, bgn: bgn + distance, :] | |
| elif self.dim == 3: | |
| input[:, :, bgn: bgn + distance] = random_sample[:, :, bgn: bgn + distance] | |
| class SpecAugmentation: | |
| def __init__(self, time_drop_width, time_stripes_num, freq_drop_width, freq_stripes_num, | |
| mask_type='mixture'): | |
| """Spec augmetation and SpecAugment++. | |
| [ref] Park, D.S., Chan, W., Zhang, Y., Chiu, C.C., Zoph, B., Cubuk, E.D. | |
| and Le, Q.V., 2019. Specaugment: A simple data augmentation method | |
| for automatic speech recognition. arXiv preprint arXiv:1904.08779. | |
| [ref] Wang H, Zou Y, Wang W., 2021. SpecAugment++: A Hidden Space | |
| Data Augmentation Method for Acoustic Scene Classification. arXiv | |
| preprint arXiv:2103.16858. | |
| Args: | |
| time_drop_width: int | |
| time_stripes_num: int | |
| freq_drop_width: int | |
| freq_stripes_num: int | |
| mask_type: str, mask type in SpecAugment++ (zero_value, mixture, cutting) | |
| """ | |
| super(SpecAugmentation, self).__init__() | |
| if mask_type == 'zero_value': | |
| self.time_augmentator = DropStripes(dim=2, drop_width=time_drop_width, | |
| stripes_num=time_stripes_num) | |
| self.freq_augmentator = DropStripes(dim=3, drop_width=freq_drop_width, | |
| stripes_num=freq_stripes_num) | |
| elif mask_type == 'mixture': | |
| self.time_augmentator = MixStripes(dim=2, mix_width=time_drop_width, | |
| stripes_num=time_stripes_num) | |
| self.freq_augmentator = MixStripes(dim=3, mix_width=freq_drop_width, | |
| stripes_num=freq_stripes_num) | |
| elif mask_type == 'cutting': | |
| self.time_augmentator = CutStripes(dim=2, cut_width=time_drop_width, | |
| stripes_num=time_stripes_num) | |
| self.freq_augmentator = CutStripes(dim=3, cut_width=freq_drop_width, | |
| stripes_num=freq_stripes_num) | |
| else: | |
| raise NameError('No such mask type in SpecAugment++') | |
| def __call__(self, inputs): | |
| # x should be in size [batch_size, channel, time_steps, freq_bins] | |
| x = self.time_augmentator(inputs) | |
| x = self.freq_augmentator(x) | |
| return x | |