From 2ac38adbe5f4e1374a079e032ed4b504351a207c Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 23 四月 2024 18:08:57 +0800
Subject: [PATCH] Dev gzf exp (#1647)

---
 funasr/bin/train.py |   29 ++++++++++++++---------------
 1 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/funasr/bin/train.py b/funasr/bin/train.py
index ab49c82..c02a66f 100644
--- a/funasr/bin/train.py
+++ b/funasr/bin/train.py
@@ -176,15 +176,12 @@
     except:
         writer = None
 
-    # if use_ddp or use_fsdp:
-    #     context = Join([model])
-    # else:
-    #     context = nullcontext()
-    context = nullcontext()
+
     for epoch in range(trainer.start_epoch, trainer.max_epoch + 1):
         time1 = time.perf_counter()
-        with context:
-            dataloader_tr, dataloader_val = dataloader.build_iter(epoch)
+        
+        for data_split_i in range(dataloader.data_split_num):
+            dataloader_tr, dataloader_val = dataloader.build_iter(epoch, data_split_i=data_split_i)
             trainer.train_epoch(
                                 model=model,
                                 optim=optim,
@@ -193,15 +190,17 @@
                                 dataloader_train=dataloader_tr,
                                 dataloader_val=dataloader_val,
                                 epoch=epoch,
-                                writer=writer
+                                writer=writer,
+                                data_split_i=data_split_i,
+                                data_split_num=dataloader.data_split_num,
                                 )
-        with context:
-            trainer.validate_epoch(
-                model=model,
-                dataloader_val=dataloader_val,
-                epoch=epoch,
-                writer=writer
-            )
+        
+        trainer.validate_epoch(
+            model=model,
+            dataloader_val=dataloader_val,
+            epoch=epoch,
+            writer=writer
+        )
         scheduler.step()
 
         

--
Gitblit v1.9.1