修复global_state没有更新的bug

light201212 · Jul 24, 2020 · 0e11758 · 0e11758
1 parent ed9737b
commit 0e11758
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tools/det_train.py b/tools/det_train.py
@@ -240,10 +240,10 @@ def train(net, optimizer, loss_func, train_loader, eval_loader, to_use_device,
                     start = time.time()
                 global_step += 1
             logger.info(f'train_loss: {train_loss / len(train_loader)}')
-            global_state['start_epoch'] = epoch
-            global_state['best_model'] = best_model
-            global_state['global_step'] = global_step
             if (epoch + 1) % train_options['val_interval'] == 0:
+                global_state['start_epoch'] = epoch
+                global_state['best_model'] = best_model
+                global_state['global_step'] = global_step
                 net_save_path = f"{train_options['checkpoint_save_dir']}/latest.pth"
                 save_checkpoint(net_save_path, net, optimizer, logger, cfg, global_state=global_state)
                 if train_options['ckpt_save_type'] == 'HighestAcc':

diff --git a/tools/rec_train.py b/tools/rec_train.py
@@ -252,6 +252,9 @@ def train(net, optimizer, scheduler, loss_func, train_loader, eval_loader, to_us
                                 f"time:{interval_batch_time:.4f}")
                     start = time.time()
                 if (i + 1) >= train_options['val_interval'] and (i + 1) % train_options['val_interval'] == 0:
+                    global_state['start_epoch'] = epoch
+                    global_state['best_model'] = best_model
+                    global_state['global_step'] = global_step
                     net_save_path = f"{train_options['checkpoint_save_dir']}/latest.pth"
                     save_checkpoint(net_save_path, net, optimizer, logger, cfg, global_state=global_state)
                     if train_options['ckpt_save_type'] == 'HighestAcc':