add my study

Author: upskyy

configs/audio/melspectrogram.yaml

@@ -0,0 +1,5 @@
+extension: pcm
+sampling_rate: 16000
+n_mel: 80
+frame_length: 20
+frame_stride: 10

configs/eval.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - audio: melspectrogram
+  - eval: default

configs/eval/default.yaml

@@ -0,0 +1,16 @@
+dataset_path: ''
+audio_path: ''
+label_path: D:/label/aihub_labels.csv
+model_path: ''
+save_transcripts_path: ''
+print_interval: 10
+num_vocabs: 2001
+pad_id: 0
+sos_id: 1
+eos_id: 2
+blank_id: 2000
+batch_size: 4
+num_workers: 4
+cuda: True
+seed: 22
+mode: eval

configs/model/deepspeech2.yaml

@@ -0,0 +1,7 @@
+architecture: deepspeech2
+input_size: 80
+hidden_size: 512
+num_layers: 3
+dropout: 0.3
+bidirectional: True
+rnn_type: gru

configs/model/joint_ctc_attention_las.yaml

@@ -0,0 +1,16 @@
+architecture: las
+input_size: 80
+encoder_hidden_size: 256
+decoder_hidden_size: 512
+encoder_layers: 3
+decoder_layers: 2
+dropout: 0.3
+bidirectional: True
+rnn_type: lstm
+teacher_forcing_ratio: 1.0
+use_joint_ctc_attention: False
+max_len: 120
+attn_mechanism: location
+smoothing: False
+ctc_weight: 0.2
+cross_entropy_weight: 0.8

configs/model/las.yaml

@@ -0,0 +1,14 @@
+architecture: las
+input_size: 80
+encoder_hidden_size: 256
+decoder_hidden_size: 512
+encoder_layers: 3
+decoder_layers: 2
+dropout: 0.3
+bidirectional: True
+rnn_type: lstm
+teacher_forcing_ratio: 1.0
+use_joint_ctc_attention: False
+max_len: 120
+attn_mechanism: location
+smoothing: False

configs/train.yaml

@@ -0,0 +1,4 @@
+defaults:
+  - audio: melspectrogram
+  - model: joint_ctc_attention_las
+  - train: las_train

configs/train/deepspeech2_train.yaml

@@ -0,0 +1,24 @@
+# Dataset
+dataset_path: D:/dataset/transcripts.txt
+audio_path: E:/KsponSpeech
+label_path: D:/label/aihub_labels.csv
+model_save_path: deepspeech2_model.pt
+
+# vocabulary
+num_vocabs: 2001
+pad_id: 0
+sos_id: 1
+eos_id: 2
+blank_id: 2000
+
+# trainer
+batch_size: 4
+num_workers: 4
+epochs: 50
+lr: 1e-06
+print_interval: 10
+
+# System
+cuda: True
+seed: 22
+mode: train

configs/train/las_train.yaml

@@ -0,0 +1,24 @@
+# Dataset
+dataset_path: D:/dataset/transcripts.txt
+audio_path: E:/KsponSpeech
+label_path: D:/label/aihub_labels.csv
+model_save_path: las_model.pt
+
+# vocabulary
+num_vocabs: 2001
+pad_id: 0
+sos_id: 1
+eos_id: 2
+blank_id: 2000
+
+# trainer
+batch_size: 4
+num_workers: 4
+epochs: 20
+lr: 1e-06
+print_interval: 10
+
+# System
+cuda: True
+seed: 22
+mode: train

eval.py

@@ -0,0 +1,69 @@
+import torch
+import numpy as np
+import random
+import hydra
+from hydra.core.config_store import ConfigStore
+from omegaconf import OmegaConf, DictConfig
+from evaluator.evaluator import Evaluator
+from model_builder import load_test_model
+from data.data_loader import (
+    SpectrogramDataset,
+    AudioDataLoader,
+)
+from vocabulary import (
+    load_label,
+    load_dataset,
+)
+
+from data import MelSpectrogramConfig
+from evaluator import EvaluateConfig
+
+
+cs = ConfigStore.instance()
+cs.store(group="audio", name="melspectrogram", node=MelSpectrogramConfig, package="audio")
+cs.store(group="eval", name="default", node=EvaluateConfig, package="eval")
+
+
+@hydra.main(config_path='configs', config_name='eval')
+def main(config: DictConfig) -> None:
+    print(OmegaConf.to_yaml(config))
+
+    torch.manual_seed(config.eval.seed)
+    torch.cuda.manual_seed_all(config.eval.seed)
+    np.random.seed(config.eval.seed)
+    random.seed(config.eval.seed)
+
+    use_cuda = config.eval.cuda and torch.cuda.is_available()
+    device = torch.device('cuda' if use_cuda else 'cpu')
+
+    char2id, id2char = load_label(config.eval.label_path, config.eval.blank_id)
+    audio_paths, transcripts, _, _ = load_dataset(config.eval.dataset_path, config.eval.mode)
+
+    test_dataset = SpectrogramDataset(
+        config.eval.audio_path,
+        audio_paths,
+        transcripts,
+        config.audio.sampling_rate,
+        config.audio.n_mel,
+        config.audio.frame_length,
+        config.audio.frame_stride,
+        config.audio.extension,
+        config.train.sos_id,
+        config.train.eos_id,
+    )
+    test_loader = AudioDataLoader(
+        test_dataset,
+        batch_size=config.eval.batch_size,
+        num_workers=config.eval.num_workers,
+    )
+
+    model = load_test_model(config, device)
+
+    print('Start Inference !!!')
+
+    evaluator = Evaluator(config, device, test_loader, id2char)
+    evaluator.evaluate(model)
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -0,0 +1,109 @@
+import torch
+import torch.optim as optim
+import numpy as np
+import random
+import hydra
+from hydra.core.config_store import ConfigStore
+from omegaconf import OmegaConf, DictConfig
+from trainer.trainer import train
+from model_builder import build_model
+
+from data.data_loader import (
+    SpectrogramDataset,
+    BucketingSampler,
+    AudioDataLoader,
+)
+from vocabulary import (
+    load_label,
+    load_dataset,
+)
+
+from data import MelSpectrogramConfig
+from models.las import (
+    ListenAttendSpellConfig,
+    JointCTCAttentionLASConfig,
+)
+from models.deepspeech2 import DeepSpeech2Config
+from trainer import (
+    ListenAttendSpellTrainConfig,
+    DeepSpeech2TrainConfig,
+)
+
+
+cs = ConfigStore.instance()
+cs.store(group="audio", name="melspectrogram", node=MelSpectrogramConfig, package="audio")
+cs.store(group="model", name="las", node=ListenAttendSpellConfig, package="model")
+cs.store(group="model", name="joint_ctc_attention_las", node=JointCTCAttentionLASConfig, package="model")
+cs.store(group="model", name="deepspeech2", node=DeepSpeech2Config, package="model")
+cs.store(group="train", name="las_train", node=ListenAttendSpellTrainConfig, package="train")
+cs.store(group="train", name="deepspeech2_train", node=DeepSpeech2TrainConfig, package="train")
+
+
+@hydra.main(config_path='configs', config_name='train')
+def main(config: DictConfig) -> None:
+    print(OmegaConf.to_yaml(config))
+
+    torch.manual_seed(config.train.seed)
+    torch.cuda.manual_seed_all(config.train.seed)
+    np.random.seed(config.train.seed)
+    random.seed(config.train.seed)
+
+    use_cuda = config.train.cuda and torch.cuda.is_available()
+    device = torch.device('cuda' if use_cuda else 'cpu')
+
+    char2id, id2char = load_label(config.train.label_path, config.train.blank_id)
+    train_audio_paths, train_transcripts, valid_audio_paths, valid_transcripts = load_dataset(config.train.dataset_path, config.train.mode)
+
+    train_dataset = SpectrogramDataset(
+        config.train.audio_path,
+        train_audio_paths,
+        train_transcripts,
+        config.audio.sampling_rate,
+        config.audio.n_mel,
+        config.audio.frame_length,
+        config.audio.frame_stride,
+        config.audio.extension,
+        config.train.sos_id,
+        config.train.eos_id,
+    )
+
+    train_sampler = BucketingSampler(train_dataset, batch_size=config.train.batch_size)
+    train_loader = AudioDataLoader(
+        train_dataset,
+        batch_sampler=train_sampler,
+        num_workers=config.train.num_workers,
+    )
+
+    valid_dataset = SpectrogramDataset(
+        config.train.audio_path,
+        valid_audio_paths,
+        valid_transcripts,
+        config.audio.sampling_rate,
+        config.audio.n_mel,
+        config.audio.frame_length,
+        config.audio.frame_stride,
+        config.audio.extension,
+        config.train.sos_id,
+        config.train.eos_id,
+    )
+    valid_sampler = BucketingSampler(valid_dataset, batch_size=config.train.batch_size)
+    valid_loader = AudioDataLoader(
+        valid_dataset,
+        batch_sampler=valid_sampler,
+        num_workers=config.train.num_workers,
+    )
+
+    model = build_model(config, device)
+    model = model.to(device)
+
+    optimizer = optim.Adam(model.parameters(), lr=config.train.lr)
+
+    print('Start Train !!!')
+    for epoch in range(0, config.train.epochs):
+        train(config, model, device, train_loader, valid_loader, train_sampler, optimizer, epoch, id2char)
+
+    torch.save(model.state_dict(), config.train.model_save_path)
+
+
+if __name__ == "__main__":
+    main()

model_builder.py

@@ -0,0 +1,74 @@
+from models.las.encoder import Encoder
+from models.las.decoder import Decoder
+from models.las.model import ListenAttendSpell
+from models.deepspeech2.model import DeepSpeech2
+from omegaconf import DictConfig
+
+import torch
+import torch.nn as nn
+
+
+def build_model(config: DictConfig, device: torch.device):
+    if config.model.architecture == 'las':
+        return build_las_model(config, device)
+
+    elif config.model.architecture == 'deepspeech2':
+        return build_ds2_model(config)
+
+
+def load_test_model(config: DictConfig, device: torch.device) -> nn.Module:
+    model = torch.load(config.eval.model_path, map_location=lambda storage, loc: storage).to(device)
+
+    model.encoder.device = device
+    model.decoder.device = device
+
+    return model
+
+
+def build_encoder(config: DictConfig) -> Encoder:
+    return Encoder(
+        config.train.num_vocabs,
+        config.model.input_size,
+        config.model.encoder_hidden_size,
+        config.model.encoder_layers,
+        config.model.dropout,
+        config.model.bidirectional,
+        config.model.rnn_type,
+        config.model.use_joint_ctc_attention
+    )
+
+
+def build_decoder(config: DictConfig, device: torch.device) -> Decoder:
+    return Decoder(
+        device,
+        config.train.num_vocabs,
+        config.model.decoder_hidden_size,
+        config.model.decoder_hidden_size,
+        config.model.decoder_layers,
+        config.model.max_len,
+        config.model.dropout,
+        config.model.rnn_type,
+        config.model.attn_mechanism,
+        config.model.smoothing,
+        config.train.sos_id,
+        config.train.eos_id,
+    )
+
+
+def build_las_model(config: DictConfig, device: torch.device) -> ListenAttendSpell:
+    encoder = build_encoder(config)
+    decoder = build_decoder(config, device)
+
+    return ListenAttendSpell(encoder, decoder)
+
+
+def build_ds2_model(config: DictConfig) -> DeepSpeech2:
+    return DeepSpeech2(
+        config.train.num_vocabs,
+        config.model.input_size,
+        config.model.hidden_size,
+        config.model.num_layers,
+        config.model.dropout,
+        config.model.bidirectional,
+        config.model.rnn_type,
+    )