nx_denoise / toolbox /torchaudio /models /rnnoise /configuration_rnnoise.py
HoneyTian's picture
update
79a2d3a
#!/usr/bin/python3
# -*- coding: utf-8 -*-
from toolbox.torchaudio.configuration_utils import PretrainedConfig
class RNNoiseConfig(PretrainedConfig):
def __init__(self,
sample_rate: int = 8000,
segment_size: int = 32000,
nfft: int = 512,
win_size: int = 512,
hop_size: int = 256,
win_type: str = "hann",
erb_bins: int = 32,
min_freq_bins_for_erb: int = 2,
conv_size: int = 128,
gru_size: int = 256,
min_snr_db: float = -10,
max_snr_db: float = 20,
lr: float = 0.001,
lr_scheduler: str = "CosineAnnealingLR",
lr_scheduler_kwargs: dict = None,
max_epochs: int = 100,
clip_grad_norm: float = 10.,
seed: int = 1234,
batch_size: int = 64,
num_workers: int = 4,
eval_steps: int = 25000,
**kwargs
):
super(RNNoiseConfig, self).__init__(**kwargs)
self.sample_rate = sample_rate
self.segment_size = segment_size
self.nfft = nfft
self.win_size = win_size
self.hop_size = hop_size
self.win_type = win_type
self.erb_bins = erb_bins
self.min_freq_bins_for_erb = min_freq_bins_for_erb
self.conv_size = conv_size
self.gru_size = gru_size
self.min_snr_db = min_snr_db
self.max_snr_db = max_snr_db
self.lr = lr
self.lr_scheduler = lr_scheduler
self.lr_scheduler_kwargs = lr_scheduler_kwargs or dict()
self.max_epochs = max_epochs
self.clip_grad_norm = clip_grad_norm
self.seed = seed
self.batch_size = batch_size
self.num_workers = num_workers
self.eval_steps = eval_steps
def main():
config = RNNoiseConfig()
config.to_yaml_file("yaml/config.yaml")
return
if __name__ == "__main__":
main()