File size: 1,123 Bytes
cae1e58 017af8f cae1e58 017af8f cae1e58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://github.com/wenet-e2e/wenet/blob/main/wenet/dataset/processor.py
"""
import random
from typing import List, Tuple
import torch
import torch.nn as nn
from torch.distributions import uniform
class SpecAugment(nn.Module):
def __init__(self,
aug_volume_factor_range: Tuple[float, float] = (0.5, 2.0),
):
super().__init__()
self.aug_volume_factor_range = aug_volume_factor_range
@staticmethod
def augment_volume(spec: torch.Tensor, factor_range: Tuple[float, float] = (0.5, 2.0)):
factor = uniform.Uniform(*factor_range)
factor = factor.sample()
spec_ = spec.clone().detach()
spec_ *= factor
return spec_
def forward(self, spec: torch.Tensor) -> torch.Tensor:
spec = self.augment_volume(spec, self.aug_volume_factor_range)
return spec
def main():
spec_augment = SpecAugment()
spec = torch.randn(size=(1, 10, 4))
print(spec)
spec_ = spec_augment.forward(spec)
print(spec_)
return
if __name__ == '__main__':
main()
|