[ssl] refine forward args to dict (#2244)

wenet-e2e · Dec 15, 2023 · 3d73821 · 3d73821
1 parent a1e1ce9
commit 3d73821
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 19 deletions.
diff --git a/wenet/ssl/bestrq/bestrq_model.py b/wenet/ssl/bestrq/bestrq_model.py
@@ -1,5 +1,5 @@
 import math
-from typing import Optional, Tuple
+from typing import Dict, Optional, Tuple
 import torch
 
 from wenet.ssl.bestrq.mask import compute_mask_indices_v2
@@ -164,11 +164,11 @@ def _reset_parameter(module: torch.nn.Module):
 
     def forward(
         self,
-        xs: torch.Tensor,
-        xs_lens: torch.Tensor,
-        text: Optional[torch.Tensor] = None,
-        text_length: Optional[torch.Tensor] = None,
+        batch: Dict,
+        device: torch.device,
     ):
+        xs = batch['feats'].to(device)
+        xs_lens = batch['feats_lengths'].to(device)
         # force global cmvn
         xs = xs - self.signal_mean
         if self.signal_norm_var:

diff --git a/wenet/ssl/w2vbert/w2vbert_model.py b/wenet/ssl/w2vbert/w2vbert_model.py
@@ -1,5 +1,5 @@
 import math
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 import torch
 
 from wenet.ssl.bestrq.mask import compute_mask_indices_v2
@@ -161,13 +161,12 @@ def _reset_parameter(module: torch.nn.Module):
     @torch.jit.ignore(drop=True)
     def forward(
         self,
-        xs: torch.Tensor,
-        xs_lens: torch.Tensor,
-        text: Optional[torch.Tensor] = None,
-        text_length: Optional[torch.Tensor] = None,
-        steps: Optional[int] = None,
+        batch: Dict,
+        device: torch.device,
     ):
-
+        steps = batch.get('steps', None)
+        xs = batch['feats'].to(device)
+        xs_lens = batch['feats_lengths'].to(device)
         assert xs.size(0) == xs_lens.size(0)
         assert steps is not None
 

diff --git a/wenet/ssl/wav2vec2/wav2vec2_model.py b/wenet/ssl/wav2vec2/wav2vec2_model.py
@@ -1,5 +1,5 @@
 import math
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 import torch
 
 import torch.nn.functional as F
@@ -220,13 +220,12 @@ def _reset_parameter(module: torch.nn.Module):
     @torch.jit.ignore(drop=True)
     def forward(
         self,
-        xs: torch.Tensor,
-        xs_lens: torch.Tensor,
-        text: Optional[torch.Tensor] = None,
-        text_length: Optional[torch.Tensor] = None,
-        steps: Optional[int] = None,
+        batch: Dict,
+        device: torch.device,
     ):
-
+        steps = batch.get('steps', None)
+        xs = batch['feats'].to(device)
+        xs_lens = batch['feats_lengths'].to(device)
         assert xs.size(0) == xs_lens.size(0)
         assert steps is not None