SWA script added

catalyst/dl/scripts/swa.py

@@ -14,13 +14,6 @@ def build_args(parser: ArgumentParser):
     parser.add_argument(
         "--models_mask", "-m", type=str, help="Pattern for models to average"
     )
-    parser.add_argument(
-        "--save-avaraged-model",
-        "-s",
-        type=bool,
-        default=True,
-        help="Flag for saving avaraged model",
-    )
 
     return parser
 
@@ -37,11 +30,8 @@ def main(args, _):
     """Main method for ``catalyst-dl swa``."""
     logdir: Path = args.logdir
     models_mask: str = args.models_mask
-    save_avaraged_model: bool = args.save_avaraged_model
 
-    averaged_weights = generate_averaged_weights(
-        logdir, models_mask, save_avaraged_model=save_avaraged_model
-    )
+    averaged_weights = generate_averaged_weights(logdir, models_mask)
 
 
 if __name__ == "__main__":

catalyst/dl/utils/swa.py

@@ -16,9 +16,9 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
     """
     Averaging of input weights.
     Args:
-        state_dicts (List[dict]): Weights to avarage
+        state_dicts (List[dict]): Weights to average
     Returns:
-        Avaraged weights
+        Averaged weights
     """
     # source https://gist.github.com/qubvel/70c3d5e4cddcde731408f478e12ef87b
 
@@ -45,17 +45,14 @@ def load_weight(path: str) -> dict:
     return weights
 
 
-def generate_averaged_weights(
-    logdir: Path, models_mask: str, save_avaraged_model: bool = True
-) -> OrderedDict:
+def generate_averaged_weights(logdir: Path, models_mask: str) -> OrderedDict:
     """
-    Averaging of input weights.
+    Averaging of input weights and saving them.
     Args:
         logdir (Path): Path to logs directory
         models_mask (str): globe-like pattern for models to average
-        save_avaraged_model (bool): Flag for saving avaraged model
     Returns:
-        Avaraged weights
+        Averaged weights
     """
 
     config_path = logdir / "configs" / "_config.json"
@@ -66,9 +63,6 @@ def generate_averaged_weights(
     all_weights = [load_weight(path) for path in models_pathes]
     averaged_dict = average_weights(all_weights)
 
-    if save_avaraged_model:
-        torch.save(
-            averaged_dict, str(logdir / "checkpoints" / "swa_weights.pth")
-        )
+    torch.save(averaged_dict, str(logdir / "checkpoints" / "swa_weights.pth"))
 
     return averaged_dict