不需要每个batch设置train()

Author: Tongjilibo

examples/tutorials_mnist.py

@@ -60,7 +60,7 @@ class MyEvaluator(Evaluator):
     # 重构评价函数
     def evaluate(self):
         total, hit = 1e-5, 0
-        for X, y in tqdm(test_dataloader, desc='Evaluating'):
+        for X, y in tqdm(test_dataloader, desc='Evaluating', ncols=80):
             pred_y = model.predict(X).argmax(dim=-1)
             hit += pred_y.eq(y).sum().item()
             total += y.shape[0]

torch4keras/trainer.py

@@ -361,7 +361,9 @@ def fit(self, train_dataloader, steps_per_epoch=None, epochs=1, callbacks=None,
                 self.callbacks.on_batch_begin(self.global_step, self.local_step, logs)
 
                 # forward和backward
-                self.unwrap_model().train()  # 设置为train模式
+                if not self.unwrap_model().training:
+                    self.unwrap_model().train()  # 设置为train模式
+                    
                 tr_loss, tr_loss_detail = 0, {}
                 for _ in range(self.grad_accumulation_steps):
                     train_X, train_y = self._prepare_nextbatch()  # 获取下一个batch的训练数据