Add segment args for pred and refine MultiSegRecord

tests/test_contrib_workflow.py

@@ -63,32 +63,46 @@ task = {
 }
 
 
-def test_multiseg():
+def train_multiseg():
     model = init_instance_by_config(task["model"])
     dataset = init_instance_by_config(task["dataset"])
     with R.start(experiment_name="workflow"):
         R.log_params(**flatten_dict(task))
         model.fit(dataset)
-
-        # prediction
         recorder = R.get_recorder()
         sr = MultiSegRecord(model, dataset, recorder)
-        sr.generate(dict(valid="valid", test="test"))
-
+        sr.generate(dict(valid="valid", test="test"), True)
         uri = R.get_uri()
+    return uri
 
+
+def train_mse():
+    model = init_instance_by_config(task["model"])
+    dataset = init_instance_by_config(task["dataset"])
+    with R.start(experiment_name="workflow"):
+        R.log_params(**flatten_dict(task))
+        model.fit(dataset)
+        recorder = R.get_recorder()
+        sr = SignalMseRecord(recorder, model=model, dataset=dataset)
+        sr.generate()
+        uri = R.get_uri()
     return uri
 
 
 class TestAllFlow(TestAutoData):
     def test_0_multiseg(self):
-        uri_path = test_multiseg()
+        uri_path = train_multiseg()
+        shutil.rmtree(str(Path(uri_path.strip("file:")).resolve()))
+
+    def test_1_mse(self):
+        uri_path = train_mse()
         shutil.rmtree(str(Path(uri_path.strip("file:")).resolve()))
 
 
 def suite():
     _suite = unittest.TestSuite()
     _suite.addTest(TestAllFlow("test_0_multiseg"))
+    _suite.addTest(TestAllFlow("test_1_mse"))
     return _suite