add fp16 and upstream states for tree pipeline (#80)

facebookresearch · Oct 11, 2023 · 5aabaf1 · 5aabaf1
1 parent b07447a
commit 5aabaf1
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/simuleval/agents/pipeline.py b/simuleval/agents/pipeline.py
@@ -269,17 +269,20 @@ def push_impl(
         module: GenericAgent,
         segment: Segment,
         states: Optional[Dict[GenericAgent, AgentStates]],
+        upstream_states: Dict[int, AgentStates],
     ):
         # DFS over the tree
         children = self.module_dict[module]
         if len(children) == 0:  # leaf node
             module.push(segment, states[module])
             return []
 
-        segment = module.pushpop(segment, states[module])
+        segment = module.pushpop(segment, states[module], upstream_states)
+        assert len(upstream_states) not in upstream_states
+        upstream_states[len(upstream_states)] = states[module]
 
         for child in children:
-            self.push_impl(child, segment, states)
+            self.push_impl(child, segment, states, upstream_states)
 
     def pushpop(
         self,
@@ -300,8 +303,11 @@ def push(
             states = {module: None for module in self.module_dict}
         else:
             assert len(states) == len(self.module_dict)
+
+        if upstream_states is None:
+            upstream_states = {}
 
-        self.push_impl(self.source_module, segment, states)
+        self.push_impl(self.source_module, segment, states, upstream_states)
 
     def pop(
         self, states: Optional[Dict[GenericAgent, AgentStates]] = None

diff --git a/simuleval/options.py b/simuleval/options.py
@@ -178,6 +178,9 @@ def general_parser():
     parser.add_argument(
         "--device", type=str, default="cpu", help="Device to run the model."
     )
+    parser.add_argument(
+        "--fp16", action="store_true", default=False, help="Use fp16."
+    )
     return parser
 
 

diff --git a/simuleval/utils/agent.py b/simuleval/utils/agent.py
@@ -141,8 +141,8 @@ def build_system_args(
 
     args = parser.parse_args(cli_argument_list(config_dict))
 
-    logger.info(f"System will run on device: {args.device}.")
-    system.to(args.device)
+    logger.info(f"System will run on device: {args.device}. fp16: {args.fp16}")
+    system.to(args.device, fp16=args.fp16)
 
     args.source_type = system.source_type
     args.target_type = system.target_type