ONNX export with NSF speedup

Author: yqzhishen

onnx/export/export_nsf_hifigan.py

@@ -7,8 +7,8 @@
 import onnxsim
 import torch
 import torch.nn as nn
-import torch.nn.functional as functional
-from torch.nn import Conv1d, ConvTranspose1d, Conv2d
+import torch.nn.functional as F
+from torch.nn import Conv1d, ConvTranspose1d
 from torch.nn.utils import weight_norm, remove_weight_norm
 
 from modules.nsf_hifigan.env import AttrDict
@@ -45,71 +45,63 @@ def __init__(self, samp_rate, harmonic_num=0,
         self.dim = self.harmonic_num + 1
         self.sampling_rate = samp_rate
         self.voiced_threshold = voiced_threshold
-        self.diff = Conv2d(
-            in_channels=1,
-            out_channels=1,
-            kernel_size=(2, 1),
-            stride=(1, 1),
-            padding=0,
-            dilation=(1, 1),
-            bias=False
-        )
-        self.diff.weight = nn.Parameter(torch.FloatTensor([[[[-1.], [1.]]]]))
 
-    def _f02sine(self, f0_values):
+    def _f02sine(self, f0_values, upp):
         """ f0_values: (batchsize, length, dim)
             where dim indicates fundamental tone and overtones
         """
-        # convert to F0 in rad. The integer part n can be ignored
-        # because 2 * np.pi * n doesn't affect phase
-        rad_values = (f0_values / self.sampling_rate).fmod(1.)
-
-        # initial phase noise (no noise for fundamental component)
+        rad_values = (f0_values / self.sampling_rate).fmod(1.)  ###%1意味着n_har的乘积无法后处理优化
         rand_ini = torch.rand(1, self.dim, device=f0_values.device)
         rand_ini[:, 0] = 0
         rad_values[:, 0, :] += rand_ini
-
-        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
-
-        # To prevent torch.cumsum numerical overflow,
-        # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
-        # Buffer tmp_over_one_idx indicates the time step to add -1.
-        # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
-        tmp_over_one = torch.cumsum(rad_values, dim=1).fmod(1.)
-
-        diff = self.diff(tmp_over_one.unsqueeze(1)).squeeze(1)  # Equivalent to torch.diff, but able to export ONNX
-        cumsum_shift = (diff < 0).float()
-        cumsum_shift = torch.cat((torch.zeros((1, 1, self.dim)).to(f0_values.device), cumsum_shift), dim=1)
-        sines = torch.sin(torch.cumsum(rad_values - cumsum_shift, dim=1) * (2 * np.pi))
+        is_half = rad_values.dtype is not torch.float32
+        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
+        if is_half:
+            tmp_over_one = tmp_over_one.half()
+        else:
+            tmp_over_one = tmp_over_one.float()
+        tmp_over_one *= upp
+        tmp_over_one = F.interpolate(
+            tmp_over_one.transpose(2, 1), scale_factor=upp,
+            mode='linear', align_corners=True
+        ).transpose(2, 1)
+        rad_values = F.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
+        tmp_over_one = tmp_over_one.fmod(1.)
+        diff = F.conv2d(
+            tmp_over_one.unsqueeze(1), torch.FloatTensor([[[[-1.], [1.]]]]).to(tmp_over_one.device),
+            stride=(1, 1), padding=0, dilation=(1, 1)
+        ).squeeze(1)  # Equivalent to torch.diff, but able to export ONNX
+        cumsum_shift = (diff < 0).double()
+        cumsum_shift = torch.cat((
+            torch.zeros((1, 1, self.dim), dtype=torch.double).to(f0_values.device),
+            cumsum_shift
+        ), dim=1)
+        sines = torch.sin(torch.cumsum(rad_values.double() + cumsum_shift, dim=1) * 2 * np.pi)
+        if is_half:
+            sines = sines.half()
+        else:
+            sines = sines.float()
         return sines
 
-    def forward(self, f0):
+
+    @torch.no_grad()
+    def forward(self, f0, upp):
         """ sine_tensor, uv = forward(f0)
         input F0: tensor(batchsize=1, length, dim=1)
                   f0 for unvoiced steps should be 0
         output sine_tensor: tensor(batchsize=1, length, dim)
         output uv: tensor(batchsize=1, length, 1)
         """
-        with torch.no_grad():
-            # fundamental component
-            fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
-
-            # generate sine waveforms
-            sine_waves = self._f02sine(fn) * self.sine_amp
-
-            # generate uv signal
-            uv = (f0 > self.voiced_threshold).float()
-
-            # noise: for unvoiced should be similar to sine_amp
-            #        std = self.sine_amp/3 -> max value ~ self.sine_amp
-            # .       for voiced regions is self.noise_std
-            noise_amp = uv * self.noise_std + (1 - uv) * (self.sine_amp / 3)
-            noise = noise_amp * torch.randn_like(sine_waves)
-
-            # first: set the unvoiced part to 0 by uv
-            # then: additive noise
-            sine_waves = sine_waves * uv + noise
-        return sine_waves, uv, noise
+        f0 = f0.unsqueeze(-1)
+        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
+        sine_waves = self._f02sine(fn, upp)
+        sine_waves = sine_waves * self.sine_amp
+        uv = (f0 > self.voiced_threshold).float()
+        uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+        sine_waves = sine_waves * uv + noise
+        return sine_waves
 
 
 class SourceModuleHnNSF(torch.nn.Module):
@@ -144,20 +136,10 @@ def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
         self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
         self.l_tanh = torch.nn.Tanh()
 
-    def forward(self, x):
-        """
-        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
-        F0_sampled (batchsize, length, 1)
-        Sine_source (batchsize, length, 1)
-        noise_source (batchsize, length 1)
-        """
-        # source for harmonic branch
-        sine_wavs, uv, _ = self.l_sin_gen(x)
+    def forward(self, x, upp):
+        sine_wavs = self.l_sin_gen(x, upp)
         sine_merge = self.l_tanh(self.l_linear(sine_wavs))
-
-        # source for noise branch, in the same shape as uv
-        noise = torch.randn_like(uv) * (self.sine_amp / 3)
-        return sine_merge, noise, uv
+        return sine_merge
 
 
 class Generator(torch.nn.Module):
@@ -166,10 +148,10 @@ def __init__(self, h):
         self.h = h
         self.num_kernels = len(h.resblock_kernel_sizes)
         self.num_upsamples = len(h.upsample_rates)
-        self.f0_upsamp = torch.nn.Upsample(scale_factor=float(np.prod(h.upsample_rates)))
         self.m_source = SourceModuleHnNSF(
             sampling_rate=h.sampling_rate,
-            harmonic_num=8)
+            harmonic_num=8
+        )
         self.noise_convs = nn.ModuleList()
         self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
         resblock = ResBlock1 if h.resblock == '1' else ResBlock2
@@ -187,24 +169,22 @@ def __init__(self, h):
             else:
                 self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
         self.resblocks = nn.ModuleList()
-        ch = None
+        ch = h.upsample_initial_channel
         for i in range(len(self.ups)):
-            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            ch //= 2
             for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
                 self.resblocks.append(resblock(h, ch, k, d))
 
         self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
         self.ups.apply(init_weights)
         self.conv_post.apply(init_weights)
+        self.upp = int(np.prod(h.upsample_rates))
 
     def forward(self, x, f0):
-        f0 = self.f0_upsamp(f0.unsqueeze(1)).transpose(1, 2)  # bs,n,t
-        har_source, noi_source, uv = self.m_source(f0)
-        har_source = har_source.transpose(1, 2)
+        har_source = self.m_source(f0, self.upp).transpose(1, 2)
         x = self.conv_pre(x)
-
         for i in range(self.num_upsamples):
-            x = functional.leaky_relu(x, LRELU_SLOPE)
+            x = F.leaky_relu(x, LRELU_SLOPE)
 
             x = self.ups[i](x)
             x_source = self.noise_convs[i](har_source)
@@ -217,10 +197,9 @@ def forward(self, x, f0):
                 else:
                     xs += self.resblocks[i * self.num_kernels + j](x)
             x = xs / self.num_kernels
-        x = functional.leaky_relu(x)
+        x = F.leaky_relu(x)
         x = self.conv_post(x)
         x = torch.tanh(x)
-
         x = x.squeeze(1)
         return x
 
@@ -261,7 +240,7 @@ def load_model(model_path, device):
     generator = Generator(h).to(device)
 
     cp_dict = torch.load(model_path)
-    generator.load_state_dict(cp_dict['generator'], strict=False)
+    generator.load_state_dict(cp_dict['generator'])
     generator.eval()
     generator.remove_weight_norm()
     del cp_dict
@@ -323,18 +302,18 @@ def export(model_path):
                     1: 'n_frames'
                 }
             },
-            opset_version=11
+            opset_version=13
         )
         print('PyTorch ONNX export finished.')
 
 
 if __name__ == '__main__':
     sys.argv = [
-        'inference/ds_e2e.py',
+        'inference/ds_cascade.py',
         '--config',
-        'configs/midi/cascade/opencs/test.yaml',
+        'configs/acoustic/nomidi.yaml',
     ]
-    path = 'onnx/assets/nsf_hifigan.onnx'
+    path = 'onnx/assets/nsf_hifigan2.onnx'
     export(path)
     simplify(path, path)
     print(f'| export \'NSF-HiFiGAN\' to \'{path}\'.')