From a9af7cd503c9daf8038153fb31c679d5fd3f5944 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Sun, 27 Sep 2020 18:45:47 +0300
Subject: [PATCH 01/14] SWA script added

---
 catalyst/dl/scripts/swa.py | 56 ++++++++++++++++++++++++++++++
 catalyst/dl/utils/swa.py   | 70 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+)
 create mode 100644 catalyst/dl/scripts/swa.py
 create mode 100644 catalyst/dl/utils/swa.py

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
new file mode 100644
index 0000000000..dddf6f9688
--- /dev/null
+++ b/catalyst/dl/scripts/swa.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+import argparse
+from argparse import ArgumentParser
+import logging
+from pathlib import Path
+
+from catalyst.dl.utils.swa import generate_averaged_weights
+
+
+def build_args(parser: ArgumentParser):
+    """Builds the command line parameters."""
+    parser.add_argument(
+        "logdir", 
+        type=Path, 
+        help="Path to models logdir"
+    )
+    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
+
+
+def parse_args():
+    """Parses the command line arguments for the main method."""
+    parser = argparse.ArgumentParser()
+    build_args(parser)
+    args = parser.parse_args()
+    return args
+
+
+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
+    )
+
+if __name__ == "__main__":
+    main(parse_args(), None)
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
new file mode 100644
index 0000000000..9e1a0ca9bd
--- /dev/null
+++ b/catalyst/dl/utils/swa.py
@@ -0,0 +1,70 @@
+import torch
+import os
+import glob
+import logging
+from typing import List
+from pathlib import Path
+from collections import OrderedDict
+from catalyst.utils import load_config
+
+logger = logging.getLogger(__name__)
+
+def average_weights(
+    state_dicts: List[dict]
+) -> OrderedDict:
+    """
+    Averaging of input weights.
+    Args:
+        state_dicts (List[dict]): Weights to avarage
+    Returns:
+        Avaraged weights
+    """
+    # source https://gist.github.com/qubvel/70c3d5e4cddcde731408f478e12ef87b
+
+    average_dict = OrderedDict()
+    for k in state_dicts[0].keys():
+        average_dict[k] = torch.true_divide(sum([state_dict[k] for state_dict in state_dicts]), len(state_dicts))
+    return average_dict
+
+def load_weight(
+    path: str
+) -> dict:
+    """
+    Load weights of a model.
+    Args:
+        path (str): Path to model weights
+    Returns:
+        Weights
+    """
+    weights = torch.load(path)
+    if "model_state_dict" in weights:
+        weights = weights["model_state_dict"]
+    return weights
+
+def generate_averaged_weights(
+    logdir: Path, 
+    models_mask: str,
+    save_avaraged_model: bool = True
+) -> OrderedDict:
+    """
+    Averaging of input weights.
+    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
+    """
+
+    config_path = logdir / "configs" / "_config.json"
+    models_pathes = glob.glob(os.path.join(logdir, "checkpoints", models_mask))
+    logging.info("Load config")
+    config: Dict[str, dict] = load_config(config_path)
+
+    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"))
+
+    return averaged_dict
\ No newline at end of file

From 27828ffd629aecf48660fd83c71e72c58d80de64 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Sun, 27 Sep 2020 18:57:29 +0300
Subject: [PATCH 02/14] Codestyle fixed

---
 catalyst/dl/scripts/swa.py | 16 ++++------------
 catalyst/dl/utils/swa.py   | 36 ++++++++++++++++++++----------------
 2 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index dddf6f9688..8b8c34fb18 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/catalyst/dl/scripts/swa.py
@@ -10,16 +10,9 @@
 
 def build_args(parser: ArgumentParser):
     """Builds the command line parameters."""
+    parser.add_argument("logdir", type=Path, help="Path to models logdir")
     parser.add_argument(
-        "logdir", 
-        type=Path, 
-        help="Path to models logdir"
-    )
-    parser.add_argument(
-        "--models_mask", 
-        "-m", 
-        type=str,
-        help="Pattern for models to average"
+        "--models_mask", "-m", type=str, help="Pattern for models to average"
     )
     parser.add_argument(
         "--save-avaraged-model",
@@ -47,10 +40,9 @@ def main(args, _):
     save_avaraged_model: bool = args.save_avaraged_model
 
     averaged_weights = generate_averaged_weights(
-        logdir,
-        models_mask,
-        save_avaraged_model=save_avaraged_model
+        logdir, models_mask, save_avaraged_model=save_avaraged_model
     )
 
+
 if __name__ == "__main__":
     main(parse_args(), None)
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 9e1a0ca9bd..3f8a436b0f 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -1,17 +1,18 @@
-import torch
-import os
+from typing import List
+from collections import OrderedDict
 import glob
 import logging
-from typing import List
+import os
 from pathlib import Path
-from collections import OrderedDict
+
+import torch
+
 from catalyst.utils import load_config
 
 logger = logging.getLogger(__name__)
 
-def average_weights(
-    state_dicts: List[dict]
-) -> OrderedDict:
+
+def average_weights(state_dicts: List[dict]) -> OrderedDict:
     """
     Averaging of input weights.
     Args:
@@ -23,12 +24,14 @@ def average_weights(
 
     average_dict = OrderedDict()
     for k in state_dicts[0].keys():
-        average_dict[k] = torch.true_divide(sum([state_dict[k] for state_dict in state_dicts]), len(state_dicts))
+        average_dict[k] = torch.true_divide(
+            sum([state_dict[k] for state_dict in state_dicts]),
+            len(state_dicts),
+        )
     return average_dict
 
-def load_weight(
-    path: str
-) -> dict:
+
+def load_weight(path: str) -> dict:
     """
     Load weights of a model.
     Args:
@@ -41,10 +44,9 @@ def load_weight(
         weights = weights["model_state_dict"]
     return weights
 
+
 def generate_averaged_weights(
-    logdir: Path, 
-    models_mask: str,
-    save_avaraged_model: bool = True
+    logdir: Path, models_mask: str, save_avaraged_model: bool = True
 ) -> OrderedDict:
     """
     Averaging of input weights.
@@ -65,6 +67,8 @@ def generate_averaged_weights(
     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
\ No newline at end of file
+    return averaged_dict

From 1f770dc1c32d68b5b79c94431c473e32caf66d36 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Mon, 28 Sep 2020 17:04:24 +0300
Subject: [PATCH 03/14] fixed pr issues

---
 catalyst/dl/scripts/swa.py | 12 +-----------
 catalyst/dl/utils/swa.py   | 18 ++++++------------
 2 files changed, 7 insertions(+), 23 deletions(-)

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index 8b8c34fb18..ca3656afc3 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/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__":
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 3f8a436b0f..8e21e807c8 100644
--- a/catalyst/dl/utils/swa.py
+++ b/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

From 78ff817fa5e67354737527b2d2770fd17a2fd9c2 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Tue, 6 Oct 2020 01:06:40 +0300
Subject: [PATCH 04/14] added test and fix

---
 catalyst/dl/tests/test_swa.py | 54 +++++++++++++++++++++++++++++++++++
 catalyst/dl/utils/swa.py      |  7 +++--
 2 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 catalyst/dl/tests/test_swa.py

diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
new file mode 100644
index 0000000000..6ba2b6e5ac
--- /dev/null
+++ b/catalyst/dl/tests/test_swa.py
@@ -0,0 +1,54 @@
+import os
+from pathlib import Path
+import shutil
+import sys
+import unittest
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from catalyst.dl.utils.swa import generate_averaged_weights
+
+sys.path.append(".")
+
+
+class Net(nn.Module):
+    def __init__(self, init_weight=4):
+        super(Net, self).__init__()
+        self.fc = nn.Linear(2, 1)
+        self.fc.weight.data.fill_(init_weight)
+        self.fc.bias.data.fill_(init_weight)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x
+
+
+class TestSwa(unittest.TestCase):
+    def setUp(self):
+        net1 = Net(init_weight=2)
+        net2 = Net(init_weight=4)
+        os.mkdir("./checkpoints")
+        torch.save(net1.state_dict(), "./checkpoints/net1.pth")
+        torch.save(net2.state_dict(), "./checkpoints/net2.pth")
+
+    def tearDown(self):
+        shutil.rmtree("./checkpoints")
+
+    def test_averaging(self):
+        generate_averaged_weights(
+            logdir=Path("./"),
+            models_mask="net*",
+            save_path=Path("./checkpoints"),
+        )
+        model = Net()
+        model.load_state_dict(torch.load("./checkpoints/swa_weights.pth"))
+
+        self.assertEqual(int(model.fc.weight.data[0][0]), 3)
+        self.assertEqual(int(model.fc.weight.data[0][1]), 3)
+        self.assertEqual(int(model.fc.bias.data[0]), 3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 8e21e807c8..7067d01d83 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -45,12 +45,15 @@ def load_weight(path: str) -> dict:
     return weights
 
 
-def generate_averaged_weights(logdir: Path, models_mask: str) -> OrderedDict:
+def generate_averaged_weights(
+    logdir: Path, models_mask: str, save_path: Path
+) -> OrderedDict:
     """
     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_path (Path): Path to save averaged model
     Returns:
         Averaged weights
     """
@@ -63,6 +66,6 @@ def generate_averaged_weights(logdir: Path, models_mask: str) -> OrderedDict:
     all_weights = [load_weight(path) for path in models_pathes]
     averaged_dict = average_weights(all_weights)
 
-    torch.save(averaged_dict, str(logdir / "checkpoints" / "swa_weights.pth"))
+    torch.save(averaged_dict, str(save_path / "swa_weights.pth"))
 
     return averaged_dict

From 8260b0a1fc6a8c94abb929b469da20f0474e90ba Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Tue, 6 Oct 2020 18:44:52 +0300
Subject: [PATCH 05/14] codestyle fix

---
 catalyst/dl/scripts/swa.py    | 17 ++++++++++++-----
 catalyst/dl/tests/test_swa.py | 23 +++++++++++------------
 catalyst/dl/utils/swa.py      | 22 ++++++++--------------
 3 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index ca3656afc3..0207c0de4c 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/catalyst/dl/scripts/swa.py
@@ -1,19 +1,21 @@
-#!/usr/bin/env python
-
 import argparse
 from argparse import ArgumentParser
-import logging
 from pathlib import Path
 
+import torch
+
 from catalyst.dl.utils.swa import generate_averaged_weights
 
 
 def build_args(parser: ArgumentParser):
     """Builds the command line parameters."""
-    parser.add_argument("logdir", type=Path, help="Path to models logdir")
+    parser.add_argument("--logdir", type=Path, help="Path to models logdir")
     parser.add_argument(
         "--models_mask", "-m", type=str, help="Pattern for models to average"
     )
+    parser.add_argument(
+        "--save_path", type=Path, help="Path to save averaged model"
+    )
 
     return parser
 
@@ -30,8 +32,13 @@ def main(args, _):
     """Main method for ``catalyst-dl swa``."""
     logdir: Path = args.logdir
     models_mask: str = args.models_mask
+    save_path: Path = args.save_path
+
+    averaged_weights = generate_averaged_weights(
+        logdir, models_mask, save_path
+    )
 
-    averaged_weights = generate_averaged_weights(logdir, models_mask)
+    torch.save(averaged_weights, str(save_path / "swa_weights.pth"))
 
 
 if __name__ == "__main__":
diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
index 6ba2b6e5ac..ddb08a059a 100644
--- a/catalyst/dl/tests/test_swa.py
+++ b/catalyst/dl/tests/test_swa.py
@@ -1,32 +1,30 @@
 import os
 from pathlib import Path
 import shutil
-import sys
 import unittest
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from catalyst.dl.utils.swa import generate_averaged_weights
 
-sys.path.append(".")
-
 
 class Net(nn.Module):
+    """Dummy network class."""
+
     def __init__(self, init_weight=4):
+        """Initialization of network and filling it with given numbers."""
         super(Net, self).__init__()
         self.fc = nn.Linear(2, 1)
         self.fc.weight.data.fill_(init_weight)
         self.fc.bias.data.fill_(init_weight)
 
-    def forward(self, x):
-        x = self.fc(x)
-        return x
-
 
 class TestSwa(unittest.TestCase):
+    """Test SWA class."""
+
     def setUp(self):
+        """Test set up."""
         net1 = Net(init_weight=2)
         net2 = Net(init_weight=4)
         os.mkdir("./checkpoints")
@@ -34,14 +32,15 @@ def setUp(self):
         torch.save(net2.state_dict(), "./checkpoints/net2.pth")
 
     def tearDown(self):
+        """Test tear down."""
         shutil.rmtree("./checkpoints")
 
     def test_averaging(self):
-        generate_averaged_weights(
-            logdir=Path("./"),
-            models_mask="net*",
-            save_path=Path("./checkpoints"),
+        """Test SWA method."""
+        weights = generate_averaged_weights(
+            logdir=Path("./"), models_mask="net*"
         )
+        torch.save(weights, str("./checkpoints/swa_weights.pth"))
         model = Net()
         model.load_state_dict(torch.load("./checkpoints/swa_weights.pth"))
 
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 7067d01d83..c288659df6 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -7,16 +7,16 @@
 
 import torch
 
-from catalyst.utils import load_config
-
 logger = logging.getLogger(__name__)
 
 
 def average_weights(state_dicts: List[dict]) -> OrderedDict:
     """
     Averaging of input weights.
+
     Args:
         state_dicts (List[dict]): Weights to average
+
     Returns:
         Averaged weights
     """
@@ -25,8 +25,7 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
     average_dict = OrderedDict()
     for k in state_dicts[0].keys():
         average_dict[k] = torch.true_divide(
-            sum([state_dict[k] for state_dict in state_dicts]),
-            len(state_dicts),
+            sum(state_dict[k] for state_dict in state_dicts), len(state_dicts),
         )
     return average_dict
 
@@ -34,8 +33,10 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
 def load_weight(path: str) -> dict:
     """
     Load weights of a model.
+
     Args:
         path (str): Path to model weights
+
     Returns:
         Weights
     """
@@ -45,27 +46,20 @@ def load_weight(path: str) -> dict:
     return weights
 
 
-def generate_averaged_weights(
-    logdir: Path, models_mask: str, save_path: Path
-) -> OrderedDict:
+def generate_averaged_weights(logdir: Path, models_mask: str) -> OrderedDict:
     """
     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_path (Path): Path to save averaged model
+
     Returns:
         Averaged weights
     """
-
-    config_path = logdir / "configs" / "_config.json"
     models_pathes = glob.glob(os.path.join(logdir, "checkpoints", models_mask))
-    logging.info("Load config")
-    config: Dict[str, dict] = load_config(config_path)
 
     all_weights = [load_weight(path) for path in models_pathes]
     averaged_dict = average_weights(all_weights)
 
-    torch.save(averaged_dict, str(save_path / "swa_weights.pth"))
-
     return averaged_dict

From 7ab43abe1c5ba242954db250701468f9624332cb Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Tue, 6 Oct 2020 23:41:43 +0300
Subject: [PATCH 06/14] fixes

---
 bin/tests/check_dl_cv.sh   |  5 +++++
 catalyst/dl/scripts/swa.py | 19 ++++++++++++-------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/bin/tests/check_dl_cv.sh b/bin/tests/check_dl_cv.sh
index 990c38a398..1e1e81b554 100755
--- a/bin/tests/check_dl_cv.sh
+++ b/bin/tests/check_dl_cv.sh
@@ -43,6 +43,11 @@ PYTHONPATH=./examples:./catalyst:${PYTHONPATH} \
   python catalyst/dl/scripts/trace.py \
   ${LOGDIR}
 
+echo 'pipeline 01 - swa'
+PYTHONPATH=./examples:./catalyst:${PYTHONPATH} \
+  python catalyst/dl/scripts/swa.py \
+  --logdir=${LOGDIR} --output-path=./swa.pth
+
 rm -rf ${LOGDIR}
 
 
diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index 0207c0de4c..69c13e94de 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/catalyst/dl/scripts/swa.py
@@ -11,10 +11,17 @@ def build_args(parser: ArgumentParser):
     """Builds the command line parameters."""
     parser.add_argument("--logdir", type=Path, help="Path to models logdir")
     parser.add_argument(
-        "--models_mask", "-m", type=str, help="Pattern for models to average"
+        "--models-mask",
+        "-m",
+        type=str,
+        default="train*",
+        help="Pattern for models to average",
     )
     parser.add_argument(
-        "--save_path", type=Path, help="Path to save averaged model"
+        "--output-path",
+        type=Path,
+        default="./swa.pth",
+        help="Path to save averaged model",
     )
 
     return parser
@@ -32,13 +39,11 @@ def main(args, _):
     """Main method for ``catalyst-dl swa``."""
     logdir: Path = args.logdir
     models_mask: str = args.models_mask
-    save_path: Path = args.save_path
+    output_path: Path = args.output_path
 
-    averaged_weights = generate_averaged_weights(
-        logdir, models_mask, save_path
-    )
+    averaged_weights = generate_averaged_weights(logdir, models_mask)
 
-    torch.save(averaged_weights, str(save_path / "swa_weights.pth"))
+    torch.save(averaged_weights, str(output_path))
 
 
 if __name__ == "__main__":

From 504f6be3aae67e47d9c516cc29e7116ec57a6409 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Wed, 7 Oct 2020 11:55:22 +0300
Subject: [PATCH 07/14] test fixed

---
 catalyst/dl/scripts/swa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index 69c13e94de..2440553e49 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/catalyst/dl/scripts/swa.py
@@ -14,7 +14,7 @@ def build_args(parser: ArgumentParser):
         "--models-mask",
         "-m",
         type=str,
-        default="train*",
+        default="*.pth",
         help="Pattern for models to average",
     )
     parser.add_argument(

From 648286972eabc55093d036e8d31b43a27611f5c8 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Wed, 7 Oct 2020 23:14:22 +0300
Subject: [PATCH 08/14] division fix

---
 catalyst/dl/tests/test_swa.py | 13 +++++++------
 catalyst/dl/utils/swa.py      |  5 +----
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
index ddb08a059a..d8f93a0cad 100644
--- a/catalyst/dl/tests/test_swa.py
+++ b/catalyst/dl/tests/test_swa.py
@@ -5,7 +5,8 @@
 
 import torch
 import torch.nn as nn
-
+import sys
+sys.path.append('.')
 from catalyst.dl.utils.swa import generate_averaged_weights
 
 
@@ -25,8 +26,8 @@ class TestSwa(unittest.TestCase):
 
     def setUp(self):
         """Test set up."""
-        net1 = Net(init_weight=2)
-        net2 = Net(init_weight=4)
+        net1 = Net(init_weight=2.)
+        net2 = Net(init_weight=5.)
         os.mkdir("./checkpoints")
         torch.save(net1.state_dict(), "./checkpoints/net1.pth")
         torch.save(net2.state_dict(), "./checkpoints/net2.pth")
@@ -44,9 +45,9 @@ def test_averaging(self):
         model = Net()
         model.load_state_dict(torch.load("./checkpoints/swa_weights.pth"))
 
-        self.assertEqual(int(model.fc.weight.data[0][0]), 3)
-        self.assertEqual(int(model.fc.weight.data[0][1]), 3)
-        self.assertEqual(int(model.fc.bias.data[0]), 3)
+        self.assertEqual(float(model.fc.weight.data[0][0]), 3.5)
+        self.assertEqual(float(model.fc.weight.data[0][1]), 3.5)
+        self.assertEqual(float(model.fc.bias.data[0]), 3.5)
 
 
 if __name__ == "__main__":
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index c288659df6..2dd016df79 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -1,14 +1,11 @@
 from typing import List
 from collections import OrderedDict
 import glob
-import logging
 import os
 from pathlib import Path
 
 import torch
 
-logger = logging.getLogger(__name__)
-
 
 def average_weights(state_dicts: List[dict]) -> OrderedDict:
     """
@@ -24,7 +21,7 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
 
     average_dict = OrderedDict()
     for k in state_dicts[0].keys():
-        average_dict[k] = torch.true_divide(
+        average_dict[k] = torch.div(
             sum(state_dict[k] for state_dict in state_dicts), len(state_dicts),
         )
     return average_dict

From 89d079c328177755489724c8eb27e7083606f17d Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Wed, 7 Oct 2020 23:37:57 +0300
Subject: [PATCH 09/14] codestyle fixed

---
 catalyst/dl/tests/test_swa.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
index d8f93a0cad..2ef447c949 100644
--- a/catalyst/dl/tests/test_swa.py
+++ b/catalyst/dl/tests/test_swa.py
@@ -1,14 +1,16 @@
 import os
 from pathlib import Path
 import shutil
+import sys
 import unittest
 
 import torch
 import torch.nn as nn
-import sys
-sys.path.append('.')
+
 from catalyst.dl.utils.swa import generate_averaged_weights
 
+sys.path.append(".")
+
 
 class Net(nn.Module):
     """Dummy network class."""
@@ -26,8 +28,8 @@ class TestSwa(unittest.TestCase):
 
     def setUp(self):
         """Test set up."""
-        net1 = Net(init_weight=2.)
-        net2 = Net(init_weight=5.)
+        net1 = Net(init_weight=2.0)
+        net2 = Net(init_weight=5.0)
         os.mkdir("./checkpoints")
         torch.save(net1.state_dict(), "./checkpoints/net1.pth")
         torch.save(net2.state_dict(), "./checkpoints/net2.pth")

From 4b0ce9909ec113050594261fcd09a816e9403c01 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Thu, 8 Oct 2020 09:57:18 +0300
Subject: [PATCH 10/14] fixed again

---
 catalyst/dl/tests/test_swa.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
index 2ef447c949..adbe5ae814 100644
--- a/catalyst/dl/tests/test_swa.py
+++ b/catalyst/dl/tests/test_swa.py
@@ -1,7 +1,6 @@
 import os
 from pathlib import Path
 import shutil
-import sys
 import unittest
 
 import torch
@@ -9,8 +8,6 @@
 
 from catalyst.dl.utils.swa import generate_averaged_weights
 
-sys.path.append(".")
-
 
 class Net(nn.Module):
     """Dummy network class."""

From a4e07c57f1c27caa9999e3d8f356c382a13af65e Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Mon, 12 Oct 2020 00:03:54 +0300
Subject: [PATCH 11/14] check keys matching added

---
 CHANGELOG.md             | 39 ++++++++++++++++++++++++++++++++++++---
 catalyst/dl/utils/swa.py | 13 +++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 56a6098045..2621de4a74 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,11 +5,43 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
-## [20.09.1] - YYYY-MM-DD
+## [YY.MM.R] - YYYY-MM-DD
+
+### Added
+
+- MRR metrics calculation ([#886](https://github.com/catalyst-team/catalyst/pull/886))
+- docs for MetricCallbacks ([#947](https://github.com/catalyst-team/catalyst/pull/947)) 
+- SoftMax, CosFace, ArcFace layers to contrib ([#939](https://github.com/catalyst-team/catalyst/pull/939))
+- ArcMargin layer to contrib ([#957](https://github.com/catalyst-team/catalyst/pull/957))
+- AdaCos to contrib ([#958](https://github.com/catalyst-team/catalyst/pull/958))
+- Manual SWA to utils (https://github.com/catalyst-team/catalyst/pull/945)
+
+### Changed
+
+- fixed path to `CHANGELOG.md` file and add information about unit test to `PULL_REQUEST_TEMPLATE.md` ([#955])(https://github.com/catalyst-team/catalyst/pull/955)
+- `catalyst-dl tune` config specification - now optuna params are grouped under `study_params` ([#947](https://github.com/catalyst-team/catalyst/pull/947))
+- `IRunner._prepare_for_stage` logic moved to `IStageBasedRunner.prepare_for_stage` ([#947](https://github.com/catalyst-team/catalyst/pull/947))
+    - now we create components in the following order: datasets/loaders, model, criterion, optimizer, scheduler, callbacks
+- `MnistMLDataset` and `MnistQGDataset` data split logic - now targets of the datasets are disjoint ([#949](https://github.com/catalyst-team/catalyst/pull/949))
+- 
+
+### Removed
+
+- 
+
+### Fixed
+
+- `AMPOptimizerCallback` - fix grad clip fn support ([#948](https://github.com/catalyst-team/catalyst/pull/948))
+- removed deprecated docs types ([#947](https://github.com/catalyst-team/catalyst/pull/947)) ([#952](https://github.com/catalyst-team/catalyst/pull/952))
+- docs for a few files ([#952](https://github.com/catalyst-team/catalyst/pull/952))
+
+
+## [20.09.1] - 2020-09-25
 
 ### Added
 
 - Runner registry support for Config API ([#936](https://github.com/catalyst-team/catalyst/pull/936))
+
 - `catalyst-dl tune` command - Optuna with Config API integration for AutoML hyperparameters optimization ([#937](https://github.com/catalyst-team/catalyst/pull/937))
 - `OptunaPruningCallback` alias for `OptunaCallback` ([#937](https://github.com/catalyst-team/catalyst/pull/937))
 - AdamP and SGDP to `catalyst.contrib.nn.criterion` ([#942](https://github.com/catalyst-team/catalyst/pull/942))
@@ -26,11 +58,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Logging double logging :) ([#936](https://github.com/catalyst-team/catalyst/pull/936))
 
+- CMCCallback ([#941](https://github.com/catalyst-team/catalyst/pull/941))
 
 ## [20.09] - 2020-09-07
 
 ### Added
 
+- `MovieLens dataset` loader ([#903](https://github.com/catalyst-team/catalyst/pull/903))
 - `force` and `bert-level` keywords to `catalyst-data text2embedding` ([#917](https://github.com/catalyst-team/catalyst/pull/917))
 - `OptunaCallback` to `catalyst.contrib` ([#915](https://github.com/catalyst-team/catalyst/pull/915))
 - `DynamicQuantizationCallback` and `catalyst-dl quantize` script for fast quantization of your model ([#890](https://github.com/catalyst-team/catalyst/pull/915))
@@ -63,7 +97,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [20.08] - 2020-08-09
 
 ### Added
-- Full metric learning pipeline including training and validation stages ([#886](https://github.com/catalyst-team/catalyst/pull/876))
 - `CMCScoreCallback` ([#880](https://github.com/catalyst-team/catalyst/pull/880))
 - kornia augmentations `BatchTransformCallback` ([#862](https://github.com/catalyst-team/catalyst/issues/862))
 - `average_precision` and `mean_average_precision` metrics ([#883](https://github.com/catalyst-team/catalyst/pull/883))
@@ -270,4 +303,4 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- 
+- 
\ No newline at end of file
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 2dd016df79..ed42885599 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -14,10 +14,23 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
     Args:
         state_dicts (List[dict]): Weights to average
 
+    Raises:
+        KeyError: If states do not match
+
     Returns:
         Averaged weights
     """
     # source https://gist.github.com/qubvel/70c3d5e4cddcde731408f478e12ef87b
+    params_keys = None
+    for i, state_dict in enumerate(state_dicts):
+        model_params_keys = list(state_dict.keys())
+        if params_keys is None:
+            params_keys = model_params_keys
+        elif params_keys != model_params_keys:
+            raise KeyError(
+                "For checkpoint {}, expected list of params: {}, "
+                "but found: {}".format(i, params_keys, model_params_keys)
+            )
 
     average_dict = OrderedDict()
     for k in state_dicts[0].keys():

From 44f750afb99aa75ec06e99f2b13a4ad9f3223833 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Mon, 12 Oct 2020 22:20:50 +0300
Subject: [PATCH 12/14] fixed test

---
 bin/tests/check_dl_cv.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bin/tests/check_dl_cv.sh b/bin/tests/check_dl_cv.sh
index d4fcec3f0d..7531cd6bca 100755
--- a/bin/tests/check_dl_cv.sh
+++ b/bin/tests/check_dl_cv.sh
@@ -44,7 +44,7 @@ PYTHONPATH=./examples:.:${PYTHONPATH} \
   ${LOGDIR}
 
 echo 'pipeline 01 - swa'
-PYTHONPATH=./examples:./catalyst:${PYTHONPATH} \
+PYTHONPATH=./examples:.:${PYTHONPATH} \
   python catalyst/dl/scripts/swa.py \
   --logdir=${LOGDIR} --output-path=./swa.pth
 

From 17a0ff1df9cd8d243c8819d3c8d0084475f19462 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Tue, 13 Oct 2020 16:40:29 +0300
Subject: [PATCH 13/14] fixed docs

---
 catalyst/dl/utils/swa.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index ed42885599..00ae87cb38 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Union
 from collections import OrderedDict
 import glob
 import os
@@ -12,7 +12,7 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
     Averaging of input weights.
 
     Args:
-        state_dicts (List[dict]): Weights to average
+        state_dicts: Weights to average
 
     Raises:
         KeyError: If states do not match
@@ -37,6 +37,7 @@ def average_weights(state_dicts: List[dict]) -> OrderedDict:
         average_dict[k] = torch.div(
             sum(state_dict[k] for state_dict in state_dicts), len(state_dicts),
         )
+
     return average_dict
 
 
@@ -45,7 +46,7 @@ def load_weight(path: str) -> dict:
     Load weights of a model.
 
     Args:
-        path (str): Path to model weights
+        path: Path to model weights
 
     Returns:
         Weights
@@ -56,13 +57,15 @@ def load_weight(path: str) -> dict:
     return weights
 
 
-def generate_averaged_weights(logdir: Path, models_mask: str) -> OrderedDict:
+def generate_averaged_weights(
+    logdir: Union[str, Path], models_mask: str
+) -> OrderedDict:
     """
     Averaging of input weights and saving them.
 
     Args:
-        logdir (Path): Path to logs directory
-        models_mask (str): globe-like pattern for models to average
+        logdir: Path to logs directory
+        models_mask: globe-like pattern for models to average
 
     Returns:
         Averaged weights

From 669bbb90ce16ace9593c907676f338655f836851 Mon Sep 17 00:00:00 2001
From: Ivan Ivashnev <ivashnev@phystech.edu>
Date: Wed, 14 Oct 2020 00:03:55 +0300
Subject: [PATCH 14/14] fix

---
 catalyst/dl/scripts/swa.py    |  4 +++-
 catalyst/dl/tests/test_swa.py |  4 ++--
 catalyst/dl/utils/swa.py      | 11 +++++++++--
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/catalyst/dl/scripts/swa.py b/catalyst/dl/scripts/swa.py
index 2440553e49..e27f8ce9a6 100644
--- a/catalyst/dl/scripts/swa.py
+++ b/catalyst/dl/scripts/swa.py
@@ -9,7 +9,9 @@
 
 def build_args(parser: ArgumentParser):
     """Builds the command line parameters."""
-    parser.add_argument("--logdir", type=Path, help="Path to models logdir")
+    parser.add_argument(
+        "--logdir", type=Path, default=None, help="Path to models logdir"
+    )
     parser.add_argument(
         "--models-mask",
         "-m",
diff --git a/catalyst/dl/tests/test_swa.py b/catalyst/dl/tests/test_swa.py
index adbe5ae814..3172e4a215 100644
--- a/catalyst/dl/tests/test_swa.py
+++ b/catalyst/dl/tests/test_swa.py
@@ -7,6 +7,7 @@
 import torch.nn as nn
 
 from catalyst.dl.utils.swa import generate_averaged_weights
+from catalyst.utils.checkpoint import load_checkpoint
 
 
 class Net(nn.Module):
@@ -42,8 +43,7 @@ def test_averaging(self):
         )
         torch.save(weights, str("./checkpoints/swa_weights.pth"))
         model = Net()
-        model.load_state_dict(torch.load("./checkpoints/swa_weights.pth"))
-
+        model.load_state_dict(load_checkpoint("./checkpoints/swa_weights.pth"))
         self.assertEqual(float(model.fc.weight.data[0][0]), 3.5)
         self.assertEqual(float(model.fc.weight.data[0][1]), 3.5)
         self.assertEqual(float(model.fc.bias.data[0]), 3.5)
diff --git a/catalyst/dl/utils/swa.py b/catalyst/dl/utils/swa.py
index 00ae87cb38..c78b37c20d 100644
--- a/catalyst/dl/utils/swa.py
+++ b/catalyst/dl/utils/swa.py
@@ -6,6 +6,8 @@
 
 import torch
 
+from catalyst.utils.checkpoint import load_checkpoint
+
 
 def average_weights(state_dicts: List[dict]) -> OrderedDict:
     """
@@ -51,7 +53,7 @@ def load_weight(path: str) -> dict:
     Returns:
         Weights
     """
-    weights = torch.load(path)
+    weights = load_checkpoint(path)
     if "model_state_dict" in weights:
         weights = weights["model_state_dict"]
     return weights
@@ -70,7 +72,12 @@ def generate_averaged_weights(
     Returns:
         Averaged weights
     """
-    models_pathes = glob.glob(os.path.join(logdir, "checkpoints", models_mask))
+    if logdir is None:
+        models_pathes = glob.glob(models_mask)
+    else:
+        models_pathes = glob.glob(
+            os.path.join(logdir, "checkpoints", models_mask)
+        )
 
     all_weights = [load_weight(path) for path in models_pathes]
     averaged_dict = average_weights(all_weights)