update

examples/sound_classification_by_lstm/run.sh

@@ -161,6 +161,11 @@ if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
   --model_dir "${file_dir}/best" \
   --serialization_dir "${file_dir}" \
 
+  python3 step_6_export_onnx_model.py \
+  --vocabulary_dir "${vocabulary_dir}" \
+  --model_dir "${file_dir}/best" \
+  --serialization_dir "${file_dir}" \
+
 fi
 
 
@@ -175,6 +180,8 @@ if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
 
   cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
 
+  cp "${file_dir}/model.onnx" "${final_model_dir}/model.onnx"
+
   cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
   cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
   cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"

examples/sound_classification_by_lstm/run_batch.sh

@@ -0,0 +1,91 @@
+#!/usr/bin/env bash
+
+
+# pretrained voicemail
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/*/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml"
+
+
+# voicemail ch64
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-en-ph-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/en-PH/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-en-sg-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/en-SG/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-en-us-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/en-US/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-es-mx-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/es-MX/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-es-pe-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/es-PE/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-fi-fi-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/fi-FI/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-id-id-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/id-ID/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-ja-jp-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/ja-JP/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-ko-kr-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/ko-KR/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-ms-my-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/ms-MY/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-pt-br-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/pt-BR/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-th-th-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/th-TH/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+
+sh run.sh --stage 0 --stop_stage 6 --system_version centos --file_folder_name file_dir --final_model_name voicemail-zh-tw-2-ch64-lstm \
+--filename_patterns "/data/tianxing/PycharmProjects/datasets/voicemail/zh-TW/wav_finished/*/*.wav" \
+--label_plan 2-voicemail \
+--config_file "yaml/lstm-classifier-2-ch64.yaml" \
+--pretrained_model "/data/tianxing/PycharmProjects/cc_audio_8/trained_models/voicemail-2-ch64-lstm.zip"
+

examples/sound_classification_by_lstm/step_3_train_model.py

@@ -27,8 +27,8 @@ from toolbox.torch.modules.loss import FocalLoss, HingeLoss, HingeLinear
 from toolbox.torch.training.metrics.categorical_accuracy import CategoricalAccuracy
 from toolbox.torch.utils.data.vocabulary import Vocabulary
 from toolbox.torch.utils.data.dataset.wave_classifier_excel_dataset import WaveClassifierExcelDataset
-from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import LSTMClassifierPretrainedModel
-from toolbox.torchaudio.models.lstm_audio_classifier.configuration_lstm_audio_classifier import LSTMClassifierConfig
+from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import WaveClassifierPretrainedModel
+from toolbox.torchaudio.models.lstm_audio_classifier.configuration_lstm_audio_classifier import WaveClassifierConfig
 
 
 def get_args():
@@ -171,7 +171,7 @@ def main():
 
     # models
     logger.info(f"prepare models. config_file: {args.config_file}")
-    config = LSTMClassifierConfig.from_pretrained(
+    config = WaveClassifierConfig.from_pretrained(
         pretrained_model_name_or_path=args.config_file,
         # num_labels=vocabulary.get_vocab_size(namespace="labels")
     )
@@ -180,7 +180,7 @@ def main():
             vocabulary.get_vocab_size(namespace="labels"),
             config.cls_head_param["num_labels"],
         ))
-    model = LSTMClassifierPretrainedModel(
+    model = WaveClassifierPretrainedModel(
         config=config,
     )
 

examples/sound_classification_by_lstm/step_4_evaluation_model.py

@@ -21,7 +21,7 @@ import torch
 from tqdm import tqdm
 
 from toolbox.torch.utils.data.vocabulary import Vocabulary
-from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import LSTMClassifierPretrainedModel
+from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import WaveClassifierPretrainedModel
 
 
 def get_args():
@@ -64,7 +64,7 @@ def main():
     logger.info("prepare vocabulary, model")
     vocabulary = Vocabulary.from_files(args.vocabulary_dir)
 
-    model = LSTMClassifierPretrainedModel.from_pretrained(
+    model = WaveClassifierPretrainedModel.from_pretrained(
         pretrained_model_name_or_path=args.model_dir,
     )
     model.to(device)

examples/sound_classification_by_lstm/step_5_export_models.py

@@ -19,7 +19,7 @@ import numpy as np
 import torch
 
 from toolbox.torch.utils.data.vocabulary import Vocabulary
-from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import LSTMClassifierPretrainedModel
+from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import WaveClassifierPretrainedModel
 
 
 def get_args():
@@ -61,7 +61,7 @@ def main():
     logger.info("prepare vocabulary, model")
     vocabulary = Vocabulary.from_files(args.vocabulary_dir)
 
-    model = LSTMClassifierPretrainedModel.from_pretrained(
+    model = WaveClassifierPretrainedModel.from_pretrained(
         pretrained_model_name_or_path=args.model_dir,
         num_labels=vocabulary.get_vocab_size(namespace="labels")
     )
@@ -99,8 +99,11 @@ def main():
     )
     script_quant_model = torch.jit.script(quantized_model)
     script_quant_model.save(serialization_dir / "script_quant_model.zip")
+
+    # onnx
+
     return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

examples/sound_classification_by_lstm/step_6_export_onnx_model.py

@@ -0,0 +1,130 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+from collections import defaultdict
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import onnxruntime as ort
+import torch
+
+from toolbox.torch.utils.data.vocabulary import Vocabulary
+from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import WaveClassifierExport
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vocabulary_dir", default="file_dir/best/vocabulary", type=str)
+    parser.add_argument("--model_dir", default="file_dir/best", type=str)
+
+    parser.add_argument("--serialization_dir", default="file_dir/best", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config():
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.DEBUG)
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(logging.Formatter(fmt))
+
+    logger = logging.getLogger(__name__)
+
+    return logger
+
+
+def main():
+    args = get_args()
+
+    serialization_dir = Path(args.serialization_dir)
+    output_file = serialization_dir / "model.onnx"
+
+    logger = logging_config()
+
+    logger.info("export models on CPU")
+    device = torch.device("cpu")
+
+    logger.info("prepare vocabulary, model")
+    vocabulary = Vocabulary.from_files(args.vocabulary_dir)
+
+    model_export = WaveClassifierExport.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+        num_labels=vocabulary.get_vocab_size(namespace="labels")
+    )
+    model_export.to(device)
+    model_export.eval()
+
+    waveform = 0 + 25 * np.random.randn(16000,)
+    waveform = np.array(waveform, dtype=np.int16)
+    waveform = waveform / (1 << 15)
+    waveform = torch.tensor(waveform, dtype=torch.float32)
+    waveform = torch.unsqueeze(waveform, dim=0)
+    waveform = waveform.to(device)
+
+    spec = model_export.wave_encoder.wave_to_mel_spectrogram(waveform) + 1e-6
+    spec = spec.log()
+    # shape = [b, f, t]
+    spec = spec.transpose(1, 2)
+    # shape = [b, t, f]
+
+    logger.info("export jit models")
+
+    inputs = spec
+
+    lstm_layer_param = model_export.config.lstm_layer_param
+    num_layers = lstm_layer_param["num_layers"]
+    hidden_size = lstm_layer_param["hidden_size"]
+    h = torch.rand(size=(num_layers, 1, hidden_size), dtype=torch.float32)
+    c = torch.rand(size=(num_layers, 1, hidden_size), dtype=torch.float32)
+
+    # onnx
+    torch.onnx.export(model_export,
+                      args=(inputs, h, c),
+                      f=output_file.as_posix(),
+                      input_names=["inputs", "h", "c"],
+                      output_names=[
+                          "logits", "new_h", "new_c",
+                      ],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size", 1: "time_steps"},
+                          "h": {1: "batch_size"},
+                          "c": {1: "batch_size"},
+                          "logits": {0: "batch_size"},
+                          "new_h": {1: "batch_size"},
+                          "new_c": {1: "batch_size"},
+                      })
+
+    ort_session = ort.InferenceSession(output_file.as_posix())
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "h": h.numpy(),
+        "c": c.numpy(),
+    }
+    output_names = [
+        "logits", "new_h", "new_c"
+    ]
+    logits, new_h, new_c = ort_session.run(output_names, input_feed)
+    print(f"logits: {logits.shape}")
+    print(f"new_h: {new_h.shape}")
+    print(f"new_c: {new_c.shape}")
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/sound_classification_by_lstm/step_8_test_onnx_model.py

@@ -0,0 +1,118 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import shutil
+import sys
+import tempfile
+import zipfile
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import onnxruntime as ort
+from scipy.io import wavfile
+import torch
+import torchaudio
+
+from project_settings import project_path
+from toolbox.torch.utils.data.vocabulary import Vocabulary
+from toolbox.torchaudio.models.lstm_audio_classifier.configuration_lstm_audio_classifier import WaveClassifierConfig
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_file",
+        # default=(project_path / "trained_models/cc_audio_8.zip").as_posix(),
+        default=(project_path / "trained_models/voicemail-ms-my-2-ch64-lstm.zip").as_posix(),
+        type=str
+    )
+    parser.add_argument(
+        "--wav_file",
+        default=r"C:\Users\tianx\Desktop\a073d03d-d280-46df-9b2d-d904965f4500_zh-CN_h3f25ivhb0c0_1719478037746.wav",
+        type=str
+    )
+
+    parser.add_argument("--device", default="cpu", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    model_file = Path(args.model_file)
+
+    device = torch.device(args.device)
+
+    with zipfile.ZipFile(model_file, "r") as f_zip:
+        out_root = Path(tempfile.gettempdir()) / "cc_audio_8"
+        print(out_root.as_posix())
+        if out_root.exists():
+            shutil.rmtree(out_root.as_posix())
+        out_root.mkdir(parents=True, exist_ok=True)
+        f_zip.extractall(path=out_root)
+
+    tgt_path = out_root / model_file.stem
+    config_file = tgt_path / "config.yaml"
+    onnx_model_file = tgt_path / "model.onnx"
+    vocab_path = tgt_path / "vocabulary"
+
+    config = WaveClassifierConfig.from_pretrained(config_file.as_posix())
+    ort_session = ort.InferenceSession(onnx_model_file.as_posix())
+    vocabulary = Vocabulary.from_files(vocab_path.as_posix())
+
+    # transform
+    wave_to_mel_spectrogram = torchaudio.transforms.MelSpectrogram(
+        sample_rate=config.mel_spectrogram_param["sample_rate"],
+        n_fft=config.mel_spectrogram_param["n_fft"],
+        win_length=config.mel_spectrogram_param["win_length"],
+        hop_length=config.mel_spectrogram_param["hop_length"],
+        f_min=config.mel_spectrogram_param["f_min"],
+        f_max=config.mel_spectrogram_param["f_max"],
+        window_fn=torch.hamming_window if config.mel_spectrogram_param["window_fn"] == "hamming" else torch.hann_window,
+        n_mels=config.mel_spectrogram_param["n_mels"],
+    )
+
+    # infer
+    sample_rate, waveform = wavfile.read(args.wav_file)
+    waveform = waveform[:16000]
+    waveform = waveform / (1 << 15)
+    waveform = torch.tensor(waveform, dtype=torch.float32)
+    waveform = torch.unsqueeze(waveform, dim=0)
+    waveform = waveform.to(device)
+
+    spec = wave_to_mel_spectrogram(waveform) + 1e-6
+    spec = spec.log()
+    # shape = [b, f, t]
+    spec = spec.transpose(1, 2)
+    # shape = [b, t, f]
+    inputs = spec
+
+    lstm_layer_param = config.lstm_layer_param
+    num_layers = lstm_layer_param["num_layers"]
+    hidden_size = lstm_layer_param["hidden_size"]
+    h = torch.zeros(size=(num_layers, 1, hidden_size), dtype=torch.float32)
+    c = torch.zeros(size=(num_layers, 1, hidden_size), dtype=torch.float32)
+
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "h": h.numpy(),
+        "c": c.numpy(),
+    }
+    output_names = [
+        "logits", "new_h", "new_c"
+    ]
+    logits, new_h, new_c = ort_session.run(output_names, input_feed)
+    print(f"logits: {logits.shape}")
+    print(f"new_h: {new_h.shape}")
+    print(f"new_c: {new_c.shape}")
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/sound_classification_by_lstm/{step_7_test_model.py → step_9_test_model.py}

@@ -16,7 +16,7 @@ import torch
 
 from project_settings import project_path
 from toolbox.torch.utils.data.vocabulary import Vocabulary
-from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import LSTMClassifierPretrainedModel
+from toolbox.torchaudio.models.lstm_audio_classifier.modeling_lstm_audio_classifier import WaveClassifierPretrainedModel
 
 
 def get_args():
@@ -58,7 +58,7 @@ def main():
 
     vocabulary = Vocabulary.from_files(vocab_path.as_posix())
 
-    model = LSTMClassifierPretrainedModel.from_pretrained(
+    model = WaveClassifierPretrainedModel.from_pretrained(
         pretrained_model_name_or_path=tgt_path.as_posix(),
     )
     model.to(device)

requirements.txt

@@ -14,3 +14,5 @@ python-dotenv
 numpy
 onnxruntime
 scipy
+onnx
+onnxruntime

toolbox/torchaudio/models/lstm_audio_classifier/configuration_lstm_audio_classifier.py

@@ -3,7 +3,7 @@
 from toolbox.torchaudio.configuration_utils import PretrainedConfig
 
 
-class LSTMClassifierConfig(PretrainedConfig):
+class WaveClassifierConfig(PretrainedConfig):
     def __init__(self,
                  mel_spectrogram_param: dict,
                  lstm_layer_param: dict,
@@ -11,7 +11,7 @@ class LSTMClassifierConfig(PretrainedConfig):
                  cls_head_param: dict,
                  **kwargs
                  ):
-        super(LSTMClassifierConfig, self).__init__(**kwargs)
+        super(WaveClassifierConfig, self).__init__(**kwargs)
         self.mel_spectrogram_param = mel_spectrogram_param
         self.lstm_layer_param = lstm_layer_param
         self.pooling_layer_param = pooling_layer_param

toolbox/torchaudio/models/lstm_audio_classifier/modeling_lstm_audio_classifier.py

@@ -7,7 +7,7 @@ import torch
 import torchaudio
 import torch.nn as nn
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE, PretrainedConfig
-from toolbox.torchaudio.models.lstm_audio_classifier.configuration_lstm_audio_classifier import LSTMClassifierConfig
+from toolbox.torchaudio.models.lstm_audio_classifier.configuration_lstm_audio_classifier import WaveClassifierConfig
 
 
 MODEL_FILE = "model.pt"
@@ -163,16 +163,18 @@ class WaveEncoder(nn.Module):
         )
 
     def forward(self, inputs: torch.Tensor):
-        # x: [batch_size, spec_dim, seq_length]
+        # x: [b, num_samples]
         x = inputs
 
         with torch.no_grad():
-            # shape = [batch_size, spec_dim, seq_length]
+            # shape = [b, f, t]
             x = self.wave_to_mel_spectrogram(x) + 1e-6
             x = x.log()
-            x = x - torch.mean(x, dim=-1, keepdim=True)
+            # x = x - torch.mean(x, dim=-1, keepdim=True)
 
+        # shape = [b, f, t]
         x = x.transpose(1, 2)
+        # shape = [b, t, f]
 
         features, h, c = self.lstm_layer.forward(x)
         # features: shape, [b, t, hidden_size]
@@ -216,13 +218,13 @@ class ClsHead(nn.Module):
         return logits
 
 
-class LSTMClassifier(nn.Module):
+class WaveClassifier(nn.Module):
     def __init__(self,
                  wave_encoder: WaveEncoder,
                  pooling_layer: PoolingLayer,
                  cls_head: ClsHead
                  ):
-        super(LSTMClassifier, self).__init__()
+        super(WaveClassifier, self).__init__()
         self.wave_encoder = wave_encoder
         self.pooling_layer = pooling_layer
         self.cls_head = cls_head
@@ -240,11 +242,11 @@ class LSTMClassifier(nn.Module):
         return logits
 
 
-class LSTMClassifierPretrainedModel(LSTMClassifier):
+class WaveClassifierPretrainedModel(WaveClassifier):
     def __init__(self,
-                 config: LSTMClassifierConfig,
+                 config: WaveClassifierConfig,
                  ):
-        super(LSTMClassifierPretrainedModel, self).__init__(
+        super(WaveClassifierPretrainedModel, self).__init__(
             wave_encoder=WaveEncoder(
                 mel_spectrogram_param=config.mel_spectrogram_param,
                 lstm_layer_param=config.lstm_layer_param,
@@ -265,7 +267,7 @@ class LSTMClassifierPretrainedModel(LSTMClassifier):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        config = LSTMClassifierConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        config = WaveClassifierConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
 
         model = cls(config)
 
@@ -301,14 +303,41 @@ class LSTMClassifierPretrainedModel(LSTMClassifier):
         return save_directory
 
 
+class WaveClassifierExport(WaveClassifierPretrainedModel):
+    def __init__(self, config: WaveClassifierConfig):
+        super(WaveClassifierExport, self).__init__(config=config)
+
+    def forward(self,
+                inputs: torch.Tensor,
+                h: torch.Tensor = None,
+                c: torch.Tensor = None,
+                ):
+        # inputs shape = [b, t, f]
+        features, h, c = self.wave_encoder.lstm_layer.forward(inputs, h=h, c=c)
+        # features: shape, [b, t, hidden_size]
+        # h: shape, [num_layers, b, hidden_size]
+        # c: shape, [num_layers, b, hidden_size]
+
+        # features shape: [b, t, f]
+        feature = self.pooling_layer.forward(features)
+        # features shape: [b, f]
+        logits = self.cls_head.forward(feature)
+        # logits shape: [batch_size, num_classes]
+        return logits, h, c
+
+
 def main():
-    config = LSTMClassifierConfig.from_pretrained("examples/lstm_classifier.yaml")
-    model = LSTMClassifierPretrainedModel(config)
-    print(model)
+    config = WaveClassifierConfig.from_pretrained("examples/lstm_classifier.yaml")
+    # model = WaveClassifierPretrainedModel(config)
+    # inputs = torch.rand(size=(1, 16000), dtype=torch.float32)
+    # outputs = model.forward(inputs)
+    # print(outputs)
+
+    model = WaveClassifierExport(config)
+    inputs = torch.rand(size=(1, 201, 80), dtype=torch.float32)
+    logits, h, c = model.export_forward(inputs)
+    print(logits)
 
-    inputs = torch.rand(size=(1, 16000), dtype=torch.float32)
-    outputs = model.forward(inputs)
-    print(outputs)
     return