Refactoring file structure

espnet_onnx/export/asr/models/__init__.py

@@ -1,29 +1,18 @@
 # decoder
 from espnet2.asr.decoder.rnn_decoder import RNNDecoder as espnetRNNDecoder
-# from espnet2.asr.encoder.conformer_encoder import \
-#     ConformerEncoder as espnetConformerEncoder
 from espnet2.asr.encoder.contextual_block_conformer_encoder import \
     ContextualBlockConformerEncoder as espnetContextualConformer
 from espnet2.asr.encoder.contextual_block_transformer_encoder import \
     ContextualBlockTransformerEncoder as espnetContextualTransformer
-# encoder
-# from espnet2.asr.encoder.rnn_encoder import RNNEncoder as espnetRNNEncoder
-# from espnet2.asr.encoder.transformer_encoder import \
-#     TransformerEncoder as espnetTransformerEncoder
-# from espnet2.asr.encoder.vgg_rnn_encoder import \
-#     VGGRNNEncoder as espnetVGGRNNEncoder
 
-from espnet_onnx.export.asr.models.ctc import CTC
+
+from espnet_onnx.export.asr.models.layers.ctc import CTC
 from espnet_onnx.export.asr.models.decoders.rnn import RNNDecoder
 from espnet_onnx.export.asr.models.decoders.transducer import TransducerDecoder
 from espnet_onnx.export.asr.models.decoders.xformer import XformerDecoder
-# from espnet_onnx.export.asr.models.encoders.conformer import ConformerEncoder
 from espnet_onnx.export.asr.models.encoders.contextual_block_xformer import \
     ContextualBlockXformerEncoder
-# from espnet_onnx.export.asr.models.encoders.rnn import RNNEncoder
-# from espnet_onnx.export.asr.models.encoders.transformer import \
-#     TransformerEncoder
-from espnet_onnx.export.asr.models.joint_network import JointNetwork
+from espnet_onnx.export.asr.models.layers.joint_network import JointNetwork
 
 try:
     from espnet2.asr.transducer.transducer_decoder import \
@@ -69,10 +58,6 @@ def get_encoder(model, frontend, preencoder, export_config, convert_map):
             export_config=export_config
         )
         return DefaultEncoder(_model, frontend, **export_config)
-    # elif isinstance(model, espnetTransformerEncoder):
-    #     return TransformerEncoder(model, frontend, preencoder, **export_config)
-    # elif isinstance(model, espnetConformerEncoder):
-    #     return ConformerEncoder(model, frontend, preencoder, **export_config)
 
 
 def get_decoder(model, export_config):

espnet_onnx/export/layers/predecoder.py → espnet_onnx/export/asr/models/decoders/predecoder.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 from espnet.nets.pytorch_backend.rnn.attentions import NoAtt
 
-from espnet_onnx.export.layers.attention import require_tanh
+from espnet_onnx.export.asr.models.layers.attention import require_tanh
 from espnet_onnx.utils.abs_model import AbsExportModel
 
 

espnet_onnx/export/asr/models/decoders/rnn.py

@@ -6,8 +6,8 @@
 import torch.nn.functional as F
 from espnet.nets.pytorch_backend.rnn.attentions import NoAtt
 
-from espnet_onnx.export.layers.attention import OnnxNoAtt, get_attention
-from espnet_onnx.export.layers.predecoder import PreDecoder
+from espnet_onnx.export.asr.models.layers.attention import OnnxNoAtt, get_attention
+from espnet_onnx.export.asr.models.decoders.predecoder import PreDecoder
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.function import make_pad_mask
 

espnet_onnx/export/asr/models/decoders/transducer.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
+from espnet_onnx.export.asr.models.layers.embed import Embedding
 from espnet_onnx.utils.abs_model import AbsExportModel
 
 

espnet_onnx/export/asr/models/decoders/xformer.py

@@ -5,9 +5,9 @@
 from espnet.nets.pytorch_backend.transformer.attention import \
     MultiHeadedAttention
 
-from espnet_onnx.export.asr.models.decoder_layer import OnnxDecoderLayer
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
-from espnet_onnx.export.asr.models.multihead_att import \
+from espnet_onnx.export.asr.models.layers.decoder_layer import OnnxDecoderLayer
+from espnet_onnx.export.asr.models.layers.embed import Embedding
+from espnet_onnx.export.asr.models.layers.multihead_att import \
     OnnxMultiHeadedAttention
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.function import subsequent_mask

espnet_onnx/export/asr/models/encoders/conformer.py

@@ -1,6 +1,4 @@
-import os
 
-import torch
 import torch.nn as nn
 from espnet.nets.pytorch_backend.transformer.attention import (
     LegacyRelPositionMultiHeadedAttention, MultiHeadedAttention,
@@ -9,11 +7,9 @@
     Conv2dSubsampling, Conv2dSubsampling2, Conv2dSubsampling6,
     Conv2dSubsampling8)
 
-from espnet_onnx.export.asr.get_config import (get_frontend_config,
-                                               get_norm_config)
-from espnet_onnx.export.asr.models.conformer_layer import OnnxConformerLayer
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
-from espnet_onnx.export.asr.models.multihead_att import (
+from espnet_onnx.export.asr.models.layers.conformer_layer import OnnxConformerLayer
+from espnet_onnx.export.asr.models.layers.embed import Embedding
+from espnet_onnx.export.asr.models.layers.multihead_att import (
     OnnxMultiHeadedAttention, OnnxRelPosMultiHeadedAttention)
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask

espnet_onnx/export/asr/models/encoders/contextual_block_xformer.py

@@ -10,7 +10,7 @@
 
 from espnet_onnx.export.asr.get_config import (get_frontend_config,
                                                get_norm_config)
-from espnet_onnx.export.asr.models.multihead_att import \
+from espnet_onnx.export.asr.models.layers.multihead_att import \
     OnnxMultiHeadedAttention
 from espnet_onnx.utils.abs_model import AbsExportModel
 

espnet_onnx/export/asr/models/encoders/rnn.py

@@ -1,17 +1,12 @@
-import os
 
 import six
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from espnet2.asr.encoder.vgg_rnn_encoder import \
-    VGGRNNEncoder as espnetVGGRNNEncoder
 from espnet.nets.pytorch_backend.rnn.encoders import RNN, RNNP, VGG2L
+from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
-from espnet_onnx.export.asr.get_config import (get_frontend_config,
-                                               get_norm_config)
 from espnet_onnx.utils.abs_model import AbsExportModel
-from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
 
 class OnnxRNNP(nn.Module):

espnet_onnx/export/asr/models/encoders/transformer.py

@@ -1,16 +1,14 @@
-import os
 
-import torch
 import torch.nn as nn
 from espnet.nets.pytorch_backend.transformer.attention import \
     MultiHeadedAttention
 from espnet.nets.pytorch_backend.transformer.subsampling import (
     Conv2dSubsampling, Conv2dSubsampling2, Conv2dSubsampling6,
     Conv2dSubsampling8)
 
-from espnet_onnx.export.asr.models.encoder_layer import OnnxEncoderLayer
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
-from espnet_onnx.export.asr.models.multihead_att import \
+from espnet_onnx.export.asr.models.layers.encoder_layer import OnnxEncoderLayer
+from espnet_onnx.export.asr.models.layers.embed import Embedding
+from espnet_onnx.export.asr.models.layers.multihead_att import \
     OnnxMultiHeadedAttention
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask

espnet_onnx/export/asr/models/frontends/s3prl.py

@@ -4,8 +4,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from espnet_onnx.export.asr.models.encoder_layer import OnnxEncoderLayer
-from espnet_onnx.export.asr.models.multihead_att import \
+from espnet_onnx.export.asr.models.layers.encoder_layer import OnnxEncoderLayer
+from espnet_onnx.export.asr.models.layers.multihead_att import \
     OnnxMultiHeadedAttention
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask

espnet_onnx/export/asr/models/language_models/transformer.py

@@ -8,9 +8,9 @@
 from espnet.nets.pytorch_backend.transformer.subsampling import (
     Conv2dSubsampling, Conv2dSubsampling6, Conv2dSubsampling8)
 
-from espnet_onnx.export.asr.models.encoder_layer import OnnxEncoderLayer
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
-from espnet_onnx.export.asr.models.multihead_att import \
+from espnet_onnx.export.asr.models.layers.encoder_layer import OnnxEncoderLayer
+from espnet_onnx.export.asr.models.layers.embed import Embedding
+from espnet_onnx.export.asr.models.layers.multihead_att import \
     OnnxMultiHeadedAttention
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask

espnet_onnx/export/asr/models/language_models/embed.py → espnet_onnx/export/asr/models/layers/embed.py

@@ -13,7 +13,7 @@
 from espnet.nets.pytorch_backend.transformer.subsampling_without_posenc import \
     Conv2dSubsamplingWOPosEnc
 
-from espnet_onnx.export.asr.models.language_models.subsampling import (
+from espnet_onnx.export.asr.models.layers.subsampling import (
     OnnxConv2dSubsampling, OnnxConv2dSubsampling2, OnnxConv2dSubsampling6,
     OnnxConv2dSubsampling8)
 

espnet_onnx/export/convert_map.yml

@@ -2,15 +2,15 @@ asr:
   # Position Embedding
   # Onnx embedding layer will use cahce the position embedding matrix
   - from: espnet.nets.pytorch_backend.transformer.embedding.LegacyRelPositionalEncoding
-    to: espnet_onnx.export.asr.models.language_models.embed.OnnxLegacyRelPositionalEncoding
+    to: espnet_onnx.export.asr.models.layers.embed.OnnxLegacyRelPositionalEncoding
   - from: espnet.nets.pytorch_backend.transformer.embedding.ScaledPositionalEncoding
-    to: espnet_onnx.export.asr.models.language_models.embed.OnnxScaledPositionalEncoding
+    to: espnet_onnx.export.asr.models.layers.embed.OnnxScaledPositionalEncoding
   - from: espnet.nets.pytorch_backend.transformer.embedding.RelPositionalEncoding
-    to: espnet_onnx.export.asr.models.language_models.embed.OnnxRelPositionalEncoding
+    to: espnet_onnx.export.asr.models.layers.embed.OnnxRelPositionalEncoding
   - from: espnet.nets.pytorch_backend.transformer.embedding.PositionalEncoding
-    to: espnet_onnx.export.asr.models.language_models.embed.OnnxPositionalEncoding
+    to: espnet_onnx.export.asr.models.layers.embed.OnnxPositionalEncoding
   - from: espnet.nets.pytorch_backend.transformer.embedding.StreamPositionalEncoding
-    to: espnet_onnx.export.asr.models.language_models.embed.OnnxStreamPositionalEncoding
+    to: espnet_onnx.export.asr.models.layers.embed.OnnxStreamPositionalEncoding
 
   # RNN encoder need to be replaced.
   - from: espnet2.asr.encoder.rnn_encoder.RNNEncoder
@@ -32,13 +32,3 @@ asr_optimization:
   - from: espnet2.asr.encoder.transformer_encoder.TransformerEncoder
     to: espnet_onnx.export.asr.models.encoders.transformer.TransformerEncoder
 
-  # Attention layer for optimization
-  # - from: espnet.nets.pytorch_backend.transformer.attention.LegacyRelPositionMultiHeadedAttention
-  #   to: espnet_onnx.export.asr.models.multihead_att.OnnxRelPosMultiHeadedAttention
-  #   is_legacy: true
-  # - from: espnet.nets.pytorch_backend.transformer.attention.RelPositionMultiHeadedAttention
-  #   to: espnet_onnx.export.asr.models.multihead_att.OnnxRelPosMultiHeadedAttention
-  #   is_legacy: false
-  # - from: espnet.nets.pytorch_backend.transformer.attention.MultiHeadedAttention
-  #   to: espnet_onnx.export.asr.models.multihead_att.OnnxMultiHeadedAttention
-

espnet_onnx/export/tts/models/tts_models/fastspeech2.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from espnet_onnx.export.asr.models.language_models.embed import Embedding
+from espnet_onnx.export.asr.models.layers.embed import Embedding
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask, normalize
 

espnet_onnx/export/tts/models/tts_models/tacotron2.py

@@ -8,8 +8,8 @@
 import torch.nn.functional as F
 from espnet.nets.pytorch_backend.rnn.attentions import AttForward, AttForwardTA
 
-from espnet_onnx.export.layers.attention import get_attention
-from espnet_onnx.export.layers.predecoder import PreDecoder
+from espnet_onnx.export.asr.models.layers.attention import get_attention
+from espnet_onnx.export.asr.models.decoders.predecoder import PreDecoder
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask, normalize
 

espnet_onnx/export/tts/models/tts_models/vits.py

@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 
-from espnet_onnx.export.asr.models.language_models.embed import get_pos_emb
+from espnet_onnx.export.asr.models.layers.embed import get_pos_emb
 from espnet_onnx.utils.abs_model import AbsExportModel
 from espnet_onnx.utils.torch_function import MakePadMask, normalize