Introduce code from Rust monorepo. This is a mess

tts/nemo/utils/export_utils.py

@@ -0,0 +1,225 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Dict, Optional, Type
+
+import onnx
+import torch
+import torch.nn as nn
+
+from nemo.utils import logging
+
+apex_available = True
+
+try:
+    from apex.normalization.fused_layer_norm import FusedLayerNorm
+
+    def replace_FusedLayerNorm(n: nn.Module) -> Optional[nn.BatchNorm2d]:
+        """
+        Replaces Apex's FusedLayerNorm with nn.LayerNorm. This is required for ONNX export.
+        Args:
+           n: the FusedLayerNorm pytorch module to replace
+        Returns:
+           Equivalent LayerNorm module
+        """
+        if not apex_available or not isinstance(n, FusedLayerNorm):
+            return None
+
+        dev = next(n.parameters()).device
+        mod = nn.LayerNorm(n.normalized_shape, eps=n.eps, elementwise_affine=n.elementwise_affine,).to(dev)
+
+        n_state = n.state_dict()
+        mod.load_state_dict(n_state)
+        return mod
+
+    default_Apex_replacements = {"FusedLayerNorm": replace_FusedLayerNorm}
+
+except Exception as e:
+    default_Apex_replacements = {}
+    apex_available = False
+
+
+def expand_Conv1D(conv1d: nn.Module) -> Optional[nn.Conv2d]:
+    """
+    Expands a Conv1D into a Conv2D. This is required for many (closed source) commercial tools with poor support for 1D Convolutions in Onnx.
+    Args:
+        conv1d: the Conv1D pytorch module to expand
+    Returns:
+        conv2d: Conv2D module with identical weights and params
+    """
+    if not isinstance(conv1d, nn.Conv1d):
+        return None
+    conv2d = nn.Conv2d(
+        conv1d.in_channels,
+        conv1d.out_channels,
+        kernel_size=(conv1d.kernel_size[0], 1),
+        stride=(conv1d.stride[0], 1),
+        padding=(conv1d.padding[0], 0),
+        dilation=(conv1d.dilation[0], 1),
+        groups=conv1d.groups,
+        padding_mode=conv1d.padding_mode,
+    ).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
+    conv2d.bias = conv1d.bias
+    conv2d.weight = nn.Parameter(conv1d.weight.unsqueeze(-1))
+    # check that expansion is valid
+    for _ in range(2):
+        sample_input = torch.rand(1, conv1d.in_channels, 256).to(
+            device=conv1d.weight.device, dtype=conv1d.weight.dtype
+        )
+        close = conv1d(sample_input).mean() - conv2d(sample_input.unsqueeze(-1)).squeeze().mean()
+        if close.abs() > 1.0:
+            raise ValueError("Unable to expand Conv1D to Conv2D")
+    return conv2d
+
+
+def expand_BatchNorm1d(bn1d: nn.Module) -> Optional[nn.BatchNorm2d]:
+    """
+    Expands a BatchNorm1d into a BatchNorm2d. This is required for many (closed source) commercial tools with poor support for BatchNorm1d in Onnx.
+    Args:
+        bn1d: the BatchNorm1d pytorch module to expand
+    Returns:
+        bn2d: BatchNorm2d module with identical weights and params
+    """
+    if not isinstance(bn1d, nn.BatchNorm1d):
+        return None
+    mod = torch.nn.BatchNorm2d(
+        bn1d.num_features,
+        eps=bn1d.eps,
+        momentum=bn1d.momentum,
+        affine=bn1d.affine,
+        track_running_stats=bn1d.track_running_stats,
+    ).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
+    bn_state = bn1d.state_dict()
+    mod.load_state_dict(bn_state)
+    return mod
+
+
+def expand_ConvTranspose1D(conv1d: nn.Module) -> Optional[nn.ConvTranspose2d]:
+    """
+    Expands a Conv1D into a Conv2D. This is required for many (closed source) commercial tools with poor support for 1D Convolutions in Onnx.
+    Args:
+        conv1d: the Conv1D pytorch module to expand
+    Returns:
+        conv2d: Conv2D module with identical weights and params
+    """
+    if not isinstance(conv1d, nn.ConvTranspose1d):
+        return None
+    conv2d = nn.ConvTranspose2d(
+        conv1d.in_channels,
+        conv1d.out_channels,
+        kernel_size=(conv1d.kernel_size[0], 1),
+        stride=(conv1d.stride[0], 1),
+        padding=(conv1d.padding[0], 0),
+        dilation=(conv1d.dilation[0], 1),
+        groups=conv1d.groups,
+        padding_mode=conv1d.padding_mode,
+    ).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
+    conv2d.bias = conv1d.bias
+    conv2d.weight = nn.Parameter(conv1d.weight.unsqueeze(-1))
+    # check that expansion is valid
+    for _ in range(2):
+        sample_input = torch.rand(1, conv1d.in_channels, 256).to(
+            device=conv1d.weight.device, dtype=conv1d.weight.dtype
+        )
+        close = conv1d(sample_input).mean() - conv2d(sample_input.unsqueeze(-1)).squeeze().mean()
+        if close.abs() > 1.0:
+            raise ValueError("Unable to expand Conv1D to Conv2D")
+    return conv2d
+
+
+def simple_replace(BaseT: Type[nn.Module], DestT: Type[nn.Module]) -> Callable[[nn.Module], Optional[nn.Module]]:
+    """
+    Generic function generator to replace BaseT module with DestT. BaseT and DestT should have same atrributes. No weights are copied.
+    Args:
+        BaseT : module type to replace
+        DestT : destination module type
+    Returns:
+        swap function to replace BaseT module with DestT
+    """
+
+    def expansion_fn(mod: nn.Module) -> Optional[nn.Module]:
+        if not isinstance(mod, BaseT):
+            return None
+        args = [getattr(mod, name, None) for name in mod.__constants__]
+        out = DestT(*args)
+        return out
+
+    return expansion_fn
+
+
+def swap_modules(model: nn.Module, mapping: Dict[str, nn.Module]):
+    """
+    This function swaps nested modules as specified by "dot paths" in mod with a desired replacement. This allows
+    for swapping nested modules through arbitrary levels if children
+
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+
+    """
+    for path, new_mod in mapping.items():
+        expanded_path = path.split(".")
+        parent_mod = model
+        for sub_path in expanded_path[:-1]:
+            parent_mod = parent_mod._modules[sub_path]  # noqa
+        parent_mod._modules[expanded_path[-1]] = new_mod  # noqa
+
+    return model
+
+
+def replace_modules(
+    model: nn.Module, expansions: Dict[str, Callable[[nn.Module], Optional[nn.Module]]] = None
+) -> nn.Module:
+    """
+    Top-level function to replace modules in model, specified by class name with a desired replacement.
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+        expansions : replacement dictionary: module class name -> replacement function generator
+    Returns:
+        model, possibly modified in-place
+    """
+    mapping: Dict[str, nn.Module] = {}
+    for name, m in model.named_modules():
+        m_type = type(m).__name__
+        if m_type in expansions:
+            swapped = expansions[m_type](m)
+            if swapped:
+                mapping[name] = swapped
+    logging.warning(f"Swapped {len(mapping)} modules")
+    swap_modules(model, mapping)
+    return model
+
+
+default_1D_2D_replacements = {
+    "Conv1d": expand_Conv1D,
+    "ConvTranspose1d": expand_ConvTranspose1D,
+    "BatchNorm1d": expand_BatchNorm1d,
+    "AdaptiveAvgPool1d": simple_replace(nn.AdaptiveAvgPool1d, nn.AdaptiveAvgPool2d),
+    "AvgPool1d": simple_replace(nn.AvgPool1d, nn.AvgPool2d),
+}
+
+
+def replace_for_export(model: nn.Module, replace_1D_2D: bool = False) -> nn.Module:
+    """
+    Top-level function to replace default set of modules in model
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+        replace_1D_2D : include 1D -> 2D replacements
+    Returns:
+        model, possibly modified in-place
+    """
+    replace_modules(model, default_Apex_replacements)
+    if replace_1D_2D:
+        # TODO: add squeeze/unsqueeze
+        replace_modules(model, default_1D_2D_replacements)