Upload utils.py

utils.py

@@ -11,7 +11,8 @@ import torch
 
 MATPLOTLIB_FLAG = False
 
-logger = logging.getLogger(__name__)
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+logger = logging
 
 
 def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
@@ -22,7 +23,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
     if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:
         optimizer.load_state_dict(checkpoint_dict['optimizer'])
     elif optimizer is None and not skip_optimizer:  
-    #else:      Disable this line if Infer and resume checkpoint,then enable the line upper
+    #else: #Disable this line if Infer ,and enable the line upper
         new_opt_dict = optimizer.state_dict()
         new_opt_dict_params = new_opt_dict['param_groups'][0]['params']
         new_opt_dict['param_groups'] = checkpoint_dict['optimizer']['param_groups']
@@ -41,12 +42,13 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
             new_state_dict[k] = saved_state_dict[k]
             assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)
         except:
-            logger.error("%s is not in the checkpoint" % k)
+            print("error, %s is not in the checkpoint" % k)
             new_state_dict[k] = v
     if hasattr(model, 'module'):
         model.module.load_state_dict(new_state_dict, strict=False)
     else:
         model.load_state_dict(new_state_dict, strict=False)
+    print("load ")
     logger.info("Loaded checkpoint '{}' (iteration {})".format(
         checkpoint_path, iteration))
     return model, optimizer, learning_rate, iteration
@@ -154,8 +156,9 @@ def get_hparams(init=True):
     parser = argparse.ArgumentParser()
     parser.add_argument('-c', '--config', type=str, default="./configs/base.json",
                         help='JSON file for configuration')
-    parser.add_argument('-m', '--model', type=str, required=True,
+    parser.add_argument('-m', '--model', type=str, default="./OUTPUT_MODEL",
                         help='Model name')
+    parser.add_argument('--cont', dest='cont', action="store_true", default=False, help="whether to continue training on the latest checkpoint")
 
     args = parser.parse_args()
     model_dir = os.path.join("./logs", args.model)
@@ -177,6 +180,7 @@ def get_hparams(init=True):
 
     hparams = HParams(**config)
     hparams.model_dir = model_dir
+    hparams.cont = args.cont
     return hparams
 
 
@@ -204,7 +208,7 @@ def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_tim
 
 def get_hparams_from_dir(model_dir):
     config_save_path = os.path.join(model_dir, "config.json")
-    with open(config_save_path, "r", encoding='utf-8') as f:
+    with open(config_save_path, "r") as f:
         data = f.read()
     config = json.loads(data)
 
@@ -214,7 +218,7 @@ def get_hparams_from_dir(model_dir):
 
 
 def get_hparams_from_file(config_path):
-    with open(config_path, "r", encoding='utf-8') as f:
+    with open(config_path, "r") as f:
         data = f.read()
     config = json.loads(data)