update

examples/sound_classification_by_lstm/step_6_export_onnx_model.py

@@ -30,6 +30,18 @@ def get_args():
 
     parser.add_argument("--serialization_dir", default="file_dir/best", type=str)
 
+    # parser.add_argument(
+    #     "--vocabulary_dir",
+    #     default=r"D:\Users\tianx\HuggingSpaces\cc_audio_8\trained_models\voicemail-zh-tw-2-ch64-lstm\voicemail-zh-tw-2-ch64-lstm\vocabulary",
+    #     type=str
+    # )
+    # parser.add_argument(
+    #     "--model_dir",
+    #     default=r"D:\Users\tianx\HuggingSpaces\cc_audio_8\trained_models\voicemail-zh-tw-2-ch64-lstm\voicemail-zh-tw-2-ch64-lstm",
+    #     type=str
+    # )
+    # parser.add_argument("--serialization_dir", default="./", type=str)
+
     args = parser.parse_args()
     return args
 
@@ -102,13 +114,17 @@ def main():
                           "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"},
-                      })
+                          "inputs": {1: "time_steps"},
+                      }
+                      # 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 = {

toolbox/torchaudio/models/cnn_audio_classifier/modeling_cnn_audio_classifier.py

@@ -308,12 +308,13 @@ class ClsHead(nn.Module):
     def forward(self, inputs: torch.Tensor):
         # inputs: [batch_size, seq_length, spec_dim]
         x = self.feedforward(inputs)
+        # x: [batch_size, seq_length, hidden_size]
 
-        # x: [batch_size, spec_dim]
         x = torch.mean(x, dim=1)
+        # x: [batch_size, hidden_size]
 
-        # logits: [batch_size, num_labels]
         logits = self.output_project_layer.forward(x)
+        # logits: [batch_size, num_labels]
         return logits