Format & polish training and webui

utils.py

@@ -1,6 +1,5 @@
 import os
 import glob
-import sys
 import argparse
 import logging
 import json
@@ -32,11 +31,13 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
         new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
         new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
         optimizer.load_state_dict(new_opt_dict)
+
     saved_state_dict = checkpoint_dict["model"]
     if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
         state_dict = model.state_dict()
+
     new_state_dict = {}
     for k, v in state_dict.items():
         try:
@@ -47,15 +48,26 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
                 v.shape,
             )
         except:
-            logger.error("%s is not in the checkpoint" % k)
+            # For upgrading from the old version
+            if "ja_bert_proj" in k:
+                v = torch.zeros_like(v)
+                logger.warn(
+                    f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
+                )
+            else:
+                logger.error(f"{k} is not in the checkpoint")
+
             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)
+
     logger.info(
         "Loaded checkpoint '{}' (iteration {})".format(checkpoint_path, iteration)
     )
+
     return model, optimizer, learning_rate, iteration
 
 
@@ -224,20 +236,33 @@ def clean_checkpoints(path_to_models="logs/44k/", n_ckpts_to_keep=2, sort_by_tim
         for f in os.listdir(path_to_models)
         if os.path.isfile(os.path.join(path_to_models, f))
     ]
-    name_key = lambda _f: int(re.compile("._(\d+)\.pth").match(_f).group(1))
-    time_key = lambda _f: os.path.getmtime(os.path.join(path_to_models, _f))
+
+    def name_key(_f):
+        return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))
+
+    def time_key(_f):
+        return os.path.getmtime(os.path.join(path_to_models, _f))
+
     sort_key = time_key if sort_by_time else name_key
-    x_sorted = lambda _x: sorted(
-        [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
-        key=sort_key,
-    )
+
+    def x_sorted(_x):
+        return sorted(
+            [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
+            key=sort_key,
+        )
+
     to_del = [
         os.path.join(path_to_models, fn)
         for fn in (x_sorted("G")[:-n_ckpts_to_keep] + x_sorted("D")[:-n_ckpts_to_keep])
     ]
-    del_info = lambda fn: logger.info(f".. Free up space by deleting ckpt {fn}")
-    del_routine = lambda x: [os.remove(x), del_info(x)]
-    rs = [del_routine(fn) for fn in to_del]
+
+    def del_info(fn):
+        return logger.info(f".. Free up space by deleting ckpt {fn}")
+
+    def del_routine(x):
+        return [os.remove(x), del_info(x)]
+
+    [del_routine(fn) for fn in to_del]
 
 
 def get_hparams_from_dir(model_dir):