replace torch.chunk with torch.split

Author: SineStriker

acoustic/dfs_models.py

@@ -16,4 +16,4 @@ def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
 class GaussianDiffusionDenoise(GaussianDiffusion):
     def forward(self, x, t, cond):
         x = self.p_sample(x, t, cond)
-        return [x, cond]
+        return x

onnx_export_singer.py

@@ -50,7 +50,7 @@ def forward(self, txt_tokens,
             mel2ph = None
         else:
             mel2ph = mel2ph[1].item()
-        
+
         if (torch.numel(txt_tokens) == 0):
             txt_tokens = None
         if (torch.numel(spk_id) == 0):
@@ -166,7 +166,7 @@ def build_model(self):
                 torch.from_numpy(np.array([0, 0]).astype(np.int64)).to(dev),
             ),
             "singer_fs.onnx",
-            # verbose=True,
+            verbose=True,
             input_names=["txt_tokens", "spk_id",
                          "pitch_midi", "midi_dur", "is_slur", "mel2ph"],
             dynamic_axes={
@@ -209,6 +209,7 @@ def build_model(self):
                 torch.rand(1, 256, 967).to(dev),
             ),
             "singer_denoise.onnx",
+            verbose=True,
             input_names=[
                 "x",
                 "t",
@@ -224,7 +225,7 @@ def build_model(self):
                     0: "batch_size",
                     1: "what",
                     2: "frames",
-                }
+                },
             },
             opset_version=11
         )

onnx_test_singer.py

@@ -142,7 +142,6 @@ def forward_model(self, inp):
                     }
                 )
                 x = torch.from_numpy(res2[0])
-                cond = torch.from_numpy(res2[1])
 
             x = x[:, 0].transpose(1, 2)
 

onnx_test_singer_gpu.py

@@ -142,7 +142,6 @@ def forward_model(self, inp):
                     }
                 )
                 x = torch.from_numpy(res2[0])
-                cond = torch.from_numpy(res2[1])
 
             x = x[:, 0].transpose(1, 2)
 

usr/diff/net.py

@@ -70,11 +70,17 @@ def forward(self, x, conditioner, diffusion_step):
 
         y = self.dilated_conv(y) + conditioner
 
-        gate, filter = torch.chunk(y, 2, dim=1)
+        # gate, filter = torch.chunk(y, 2, dim=1)
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        gate, filter = torch.split(y, torch.div(y.shape[1], 2), dim=1)
+
         y = torch.sigmoid(gate) * torch.tanh(filter)
 
         y = self.output_projection(y)
-        residual, skip = torch.chunk(y, 2, dim=1)
+        # residual, skip = torch.chunk(y, 2, dim=1)
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        residual, skip = torch.split(y, torch.div(y.shape[1], 2), dim=1)
+        
         return (x + residual) / sqrt(2.0), skip