fix infer

heqi1990 · Sep 20, 2023 · f87d065 · f87d065
1 parent 16bd307
commit f87d065
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 14 deletions.
diff --git a/inference/me_infer.py b/inference/me_infer.py
@@ -64,11 +64,10 @@ def preprocess(self, waveform: np.ndarray) -> Dict[str, torch.Tensor]:
 
     @torch.no_grad()
     def forward_model(self, sample: Dict[str, torch.Tensor]):
-        sig=False
-        if self.config['use_BCEWithLogitsLoss']:
-            sig=True
 
-        probs, bounds = self.model(x=sample['units'], f0=sample['pitch'], mask=sample['masks'],sig=sig)
+
+
+        probs, bounds = self.model(x=sample['units'], f0=sample['pitch'], mask=sample['masks'],sig=True)
 
         return {
             'probs': probs,

diff --git a/modules/conform/Gconform.py b/modules/conform/Gconform.py
@@ -137,6 +137,6 @@ def forward(self, x, pitch, mask=None):
         midiout = self.outln(x)
         cutprp = torch.sigmoid(cutprp)
         cutprp = torch.squeeze(cutprp, -1)
-        if  self.sig:
-            midiout = torch.sigmoid(midiout)
+        # if  self.sig:
+        #     midiout = torch.sigmoid(midiout)
         return midiout, cutprp
diff --git a/training/me_task.py b/training/me_task.py
@@ -70,10 +70,8 @@ def build_model(self):
         return model
 
     def build_losses_and_metrics(self):
-        if not self.config['use_BCEWithLogitsLoss']:
-            self.midi_loss = nn.BCELoss()
-        else:
-            self.midi_loss = nn.BCEWithLogitsLoss()
+
+        self.midi_loss = nn.BCEWithLogitsLoss()
         self.bound_loss = modules.losses.BinaryEMDLoss()
         # self.bound_loss = modules.losses.BinaryEMDLoss(bidirectional=True)
         self.register_metric('midi_acc', modules.metrics.MIDIAccuracy(tolerance=0.5))
@@ -90,16 +88,16 @@ def run_model(self, sample, infer=False):
         # mask=None
 
         f0 = sample['pitch']
-        sig=False
-        if self.config['use_BCEWithLogitsLoss'] and infer:
-            sig=True
 
-        probs, bounds = self.model(x=spec, f0=f0, mask=mask,sig=sig)
+
+
 
         if infer:
+            probs, bounds = self.model(x=spec, f0=f0, mask=mask, sig=True)
             return probs, bounds
         else:
             losses = {}
+            probs, bounds = self.model(x=spec, f0=f0, mask=mask, sig=False)
 
             if  self.cfg['use_bound_loss']:
                 bound_loss = self.bound_loss(bounds, sample['bounds'])