更新 'train_for_c2net.py'

2023-06-12 10:55:50 +08:00 · 2023-06-12 10:55:50 +08:00 · 7aab8b05c6
parent 81273b6c2c
commit 7aab8b05c6
1 changed files with 8 additions and 5 deletions
--- a/train_for_c2net.py
+++ b/train_for_c2net.py
@ -112,10 +112,13 @@ if __name__ == "__main__":
    if (args.epoch_size):
        epoch_size = args.epoch_size
        print('epoch_size is: ', epoch_size)
+    model.train(epoch_size, ds_train,callbacks=[time_cb, ckpoint_cb,LossMonitor()])
    # set callback functions
-    callback =[time_cb,LossMonitor()]
-    local_rank=int(os.getenv('RANK_ID'))
+    # callback =[time_cb,LossMonitor()]
+    # local_rank=int(os.getenv('RANK_ID'))
    # for data parallel, only save checkpoint on rank 0
-    if local_rank==0 :
-        callback.append(ckpoint_cb) 
-    model.train(epoch_size,ds_train,callbacks=callback)
+    # if local_rank==0 :
+    #     callback.append(ckpoint_cb) 
+    # model.train(epoch_size,
+    #             ds_train,
+    #             callbacks=callback)