update

examples/sound_classification_by_lstm/step_6_export_onnx_model.py

@@ -14,14 +14,13 @@ 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 WaveClassifierPreprocessExport, WaveClassifierExport
+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("--onnx_preprocess_file", default="preprocess.onnx", type=str)
     parser.add_argument("--onnx_model_file", default="model.onnx", type=str)
 
     args = parser.parse_args()
@@ -46,7 +45,6 @@ def logging_config():
 def main():
     args = get_args()
 
-    onnx_preprocess_file = Path(args.onnx_preprocess_file)
     onnx_model_file = Path(args.onnx_model_file)
 
     logger = logging_config()
@@ -57,13 +55,6 @@ def main():
     logger.info("prepare vocabulary, model")
     vocabulary = Vocabulary.from_files(args.vocabulary_dir)
 
-    model_preprocess_export = WaveClassifierPreprocessExport.from_pretrained(
-        pretrained_model_name_or_path=args.model_dir,
-        num_labels=vocabulary.get_vocab_size(namespace="labels")
-    )
-    model_preprocess_export.to(device)
-    model_preprocess_export.eval()
-
     model_export = WaveClassifierExport.from_pretrained(
         pretrained_model_name_or_path=args.model_dir,
         num_labels=vocabulary.get_vocab_size(namespace="labels")
@@ -78,33 +69,8 @@ def main():
     waveform = torch.unsqueeze(waveform, dim=0)
     waveform = waveform.to(device)
 
-    logger.info("export onnx preprocess models")
-    torch.onnx.export(model_preprocess_export,
-                      args=(waveform,),
-                      f=onnx_preprocess_file.as_posix(),
-                      input_names=["inputs"],
-                      output_names=["spec"],
-                      dynamic_axes={
-                          "inputs": {1: "num_samples"},
-                      }
-                      )
-
-    preprocess_ort_session = ort.InferenceSession(onnx_preprocess_file.as_posix())
-    input_feed = {
-        "inputs": waveform.numpy(),
-    }
-    output_names = [
-        "spec",
-    ]
-    outputs = preprocess_ort_session.run(output_names, input_feed)
-    spec = outputs[0]
-    # shape = [b, t, f]
-
     logger.info("export onnx models")
 
-    inputs = spec
-    inputs = torch.tensor(inputs, dtype=torch.float32)
-
     lstm_layer_param = model_export.config.lstm_layer_param
     num_layers = lstm_layer_param["num_layers"]
     hidden_size = lstm_layer_param["hidden_size"]
@@ -112,7 +78,7 @@ def main():
     c = torch.rand(size=(num_layers, 1, hidden_size), dtype=torch.float32)
 
     torch.onnx.export(model_export,
-                      args=(inputs, h, c),
+                      args=(waveform, h, c),
                       f=onnx_model_file.as_posix(),
                       input_names=["inputs", "h", "c"],
                       output_names=[
@@ -125,7 +91,7 @@ def main():
 
     model_ort_session = ort.InferenceSession(onnx_model_file.as_posix())
     input_feed = {
-        "inputs": inputs.numpy(),
+        "inputs": waveform.numpy(),
         "h": h.numpy(),
         "c": c.numpy(),
     }

examples/sound_classification_by_lstm/step_9_evaluation_onnx_model.py

@@ -40,7 +40,6 @@ def get_args():
     parser.add_argument("--dataset", default="evaluation.xlsx", type=str)
     parser.add_argument("--vocabulary_dir", default="vocabulary", type=str)
     parser.add_argument("--model_dir", default="best", type=str)
-    parser.add_argument("--onnx_preprocess_file", default="preprocess.onnx", type=str)
     parser.add_argument("--onnx_model_file", default="model.onnx", type=str)
     parser.add_argument("--output_file", default="evaluation_onnx.xlsx", type=str)
     # parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", type=str)
@@ -70,23 +69,17 @@ def main():
             f_zip.extractall(path=out_root)
         tgt_path = out_root / model_file.stem
         config_file = tgt_path / "config.yaml"
-        onnx_preprocess_file = tgt_path / "preprocess.onnx"
         onnx_model_file = tgt_path / "model.onnx"
         vocab_path = tgt_path / "vocabulary"
         evaluation_file = tgt_path / "evaluation.xlsx"
     else:
         config_file = model_dir / "config.yaml"
-        onnx_preprocess_file = Path(args.onnx_preprocess_file)
         onnx_model_file = Path(args.onnx_model_file)
         vocab_path = Path(args.vocabulary_dir)
         evaluation_file = Path(args.dataset)
 
     config = WaveClassifierConfig.from_pretrained(config_file.as_posix())
-    preprocess_ort_session = ort.InferenceSession(
-        onnx_preprocess_file.as_posix(),
-        providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
-    )
-    model_ort_session = ort.InferenceSession(
+    ort_session = ort.InferenceSession(
         onnx_model_file.as_posix(),
         providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
     )
@@ -132,25 +125,16 @@ def main():
             end = begin + 4000
 
             waveform_ = waveform[:, begin: end]
-            input_feed = {
-                "inputs": waveform_.numpy(),
-            }
-            output_names = [
-                "spec",
-            ]
-            outputs = preprocess_ort_session.run(output_names, input_feed)
-            # shape = [b, t, f]
-            inputs: np.ndarray = outputs[0]
 
             input_feed = {
-                "inputs": inputs,
+                "inputs": waveform_.numpy(),
                 "h": h,
                 "c": c,
             }
             output_names = [
                 "logits", "new_h", "new_c"
             ]
-            logits, new_h, new_c = model_ort_session.run(output_names, input_feed)
+            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}")

toolbox/torchaudio/models/lstm_audio_classifier/modeling_lstm_audio_classifier.py

@@ -308,12 +308,14 @@ class WaveClassifierPretrainedModel(WaveClassifier):
         return save_directory
 
 
-class WaveClassifierPreprocessExport(WaveClassifierPretrainedModel):
+class WaveClassifierExport(WaveClassifierPretrainedModel):
     def __init__(self, config: WaveClassifierConfig):
-        super(WaveClassifierPreprocessExport, self).__init__(config=config)
+        super(WaveClassifierExport, self).__init__(config=config)
 
     def forward(self,
                 inputs: torch.Tensor,
+                h: torch.Tensor = None,
+                c: torch.Tensor = None,
                 ):
         # x: [b, num_samples]
         x = inputs
@@ -327,22 +329,8 @@ class WaveClassifierPreprocessExport(WaveClassifierPretrainedModel):
             # shape = [b, t, mel_bins]
             spec = x + 1e-6
             spec = spec.log()
-
-        # shape = [b, t, f]
-        return spec
-
-
-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)
+        # spec shape = [b, t, f]
+        features, h, c = self.wave_encoder.lstm_layer.forward(spec, h=h, c=c)
         # features: shape, [b, t, hidden_size]
         # h: shape, [num_layers, b, hidden_size]
         # c: shape, [num_layers, b, hidden_size]
@@ -358,10 +346,8 @@ class WaveClassifierExport(WaveClassifierPretrainedModel):
 def main():
     config = WaveClassifierConfig.from_pretrained("examples/lstm_classifier.yaml")
     model = WaveClassifierPretrainedModel(config)
-    model_preprocess = WaveClassifierPreprocessExport(config)
     model_export = WaveClassifierExport(config)
     model.eval()
-    model_preprocess.eval()
     model_export.eval()
 
     inputs = torch.rand(size=(1, 16000), dtype=torch.float32)
@@ -369,8 +355,7 @@ def main():
     logits = model.forward(inputs)
     print(logits)
 
-    spec = model_preprocess.forward(inputs)
-    logits, h, c = model_export.forward(spec)
+    logits, h, c = model_export.forward(inputs)
 
     return
 

voicemail-es-mx-2-l3-ch64-lstm.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4875557ee77aa194cba21c84bb577383ae5d6aab53a424f1253bb84e2253049e
+size 4528606