Feature: optimizer.zero_grad() speed up option

bin/tests/check_dl_core_callbacks.sh

@@ -822,5 +822,31 @@ check_line_counts ${EXP_OUTPUT} "loaded state .*/last_full.pth (global epoch 2,
 
 rm -rf ./tests/logs/_tests_dl_callbacks ${EXP_OUTPUT}
 
+################################  pipeline 23  ################################
+# checking optimizer use_fast_zero_grad
+LOG_MSG='pipeline 23'
+echo ${LOG_MSG}
+
+PYTHONPATH=./examples:./catalyst:${PYTHONPATH} \
+  python catalyst/dl/scripts/run.py \
+  --stages/stage1/stage_params/num_epochs='2:int' \
+  --stages/stage1/callbacks_params/optimizer/use_fast_zero_grad='1:bool' \
+  --expdir=${EXPDIR} \
+  --config=${EXPDIR}/config0.yml \
+  --logdir=${LOGDIR} > ${EXP_OUTPUT}
+
+cat ${EXP_OUTPUT}
+check_line_counts ${EXP_OUTPUT} "=> Loading" 0
+
+check_file_existence ${LOGFILE}
+cat ${LOGFILE}
+echo ${LOG_MSG}
+
+check_checkpoints "${CHECKPOINTS}/best" 1
+check_checkpoints "${CHECKPOINTS}/last" 1
+check_checkpoints "${CHECKPOINTS}/stage1\.[[:digit:]]" 1
+check_num_files ${CHECKPOINTS} 7   # 3x2 checkpoints + metrics.json
+
+rm -rf ${LOGDIR} ${EXP_OUTPUT}
 
 rm -rf {LOGDIR}

catalyst/core/callbacks/optimizer.py

@@ -27,6 +27,7 @@ def __init__(
         grad_clip_params: Dict = None,
         decouple_weight_decay: bool = True,
         loss_key: str = None,
+        use_fast_zero_grad: bool = False,
         xla_barrier: bool = True,
     ):
         """
@@ -37,8 +38,10 @@ def __init__(
             accumulation_steps (int): number of steps before
                 ``model.zero_grad()``
             grad_clip_params (dict): params for gradient clipping
-            decouple_weight_decay (bool): If True - decouple weight decay
+            decouple_weight_decay (bool): If ``True`` - decouple weight decay
                 regularization.
+            use_fast_zero_grad (bool): boost ``optiomizer.zero_grad()``,
+                default is ``False``.
             xla_barrier (bool): barrier option for xla. Here you can find
                 more about usage of `barrier flag
                 <https://pytorch.org/xla/release/1.5/index.html?
@@ -72,6 +75,7 @@ def __init__(
         self._optimizer_wd: List[float] = [0.0]
         self._optimizer_step_fn: Callable = None
         self.is_xla = False
+        self.use_fast_zero_grad = use_fast_zero_grad
         self.use_xla_barrier = xla_barrier
 
     def grad_step(
@@ -217,8 +221,17 @@ def on_batch_end(self, runner: IRunner) -> None:
                 optimizer_wds=self._optimizer_wd,
                 grad_clip_fn=self.grad_clip_fn,
             )
+            if not self.use_fast_zero_grad:
+                utils.maybe_recursive_call(self._optimizer, "zero_grad")
+            else:
 
-            utils.maybe_recursive_call(self._optimizer, "zero_grad")
+                def zero_grad(optimizer):
+                    for group in optimizer.param_groups:
+                        for p in group["params"]:
+                            if p.grad is not None:
+                                p.grad = None
+
+                utils.maybe_recursive_call(self._optimizer, zero_grad)
             self._accumulation_counter = 0
 
 

catalyst/core/tests/test_optimizer_callback.py

@@ -0,0 +1,80 @@
+# flake8: noqa
+import random
+
+import torch
+import torch.nn as nn
+
+from catalyst.core.callbacks import OptimizerCallback
+
+
+class DummyRunner:
+    def __init__(
+        self, loss_value: torch.tensor, optimizer: torch.optim.Optimizer
+    ):
+        self.batch_metrics = {"loss": loss_value}
+        self.is_train_loader = True
+        self.optimizer = optimizer
+        self.device = torch.device("cpu")
+
+    def get_attr(self, key, *args, **kwargs):
+        return getattr(self, key)
+
+
+def test_zero_grad():
+    model = nn.Linear(10, 2)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.BCEWithLogitsLoss()
+
+    batch_size = 3
+    inp = torch.randn(batch_size, 10)
+    target = torch.FloatTensor(batch_size, 2).uniform_()
+
+    callback = OptimizerCallback(metric_key="loss", use_fast_zero_grad=False)
+
+    loss1 = criterion(model(inp), target)
+    loss1_value = loss1.detach().item()
+
+    runner = DummyRunner(loss1, optimizer)
+
+    callback.on_stage_start(runner)
+    callback.on_epoch_start(runner)
+    callback.on_batch_end(runner)
+
+    loss2 = criterion(model(inp), target)
+    loss2_value = loss2.detach().item()
+
+    runner.batch_metrics = {"loss": loss2}
+    callback.on_epoch_start(runner)
+    callback.on_batch_end(runner)
+
+    assert loss1_value > loss2_value
+
+
+def test_fast_zero_grad():
+    model = nn.Linear(10, 2)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.BCEWithLogitsLoss()
+
+    batch_size = 3
+    inp = torch.randn(batch_size, 10)
+    target = torch.FloatTensor(batch_size, 2).uniform_()
+
+    callback = OptimizerCallback(metric_key="loss", use_fast_zero_grad=True)
+
+    loss1 = criterion(model(inp), target)
+    loss1_value = loss1.detach().item()
+
+    runner = DummyRunner(loss1, optimizer)
+
+    callback.on_stage_start(runner)
+    callback.on_epoch_start(runner)
+    callback.on_batch_end(runner)
+
+    loss2 = criterion(model(inp), target)
+    loss2_value = loss2.detach().item()
+
+    runner.batch_metrics = {"loss": loss2}
+    callback.on_epoch_start(runner)
+    callback.on_batch_end(runner)
+
+    assert loss1_value > loss2_value

catalyst/utils/misc.py

@@ -1,6 +1,6 @@
 # flake8: noqa
 # @TODO: code formatting issue for 20.07 release
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Union
 from datetime import datetime
 import inspect
 from pathlib import Path
@@ -9,7 +9,7 @@
 
 def maybe_recursive_call(
     object_or_dict,
-    method: str,
+    method: Union[str, Callable],
     recursive_args=None,
     recursive_kwargs=None,
     **kwargs,
@@ -46,7 +46,10 @@ def maybe_recursive_call(
     if not isinstance(r_args, (list, tuple)):
         r_args = [r_args]
     r_kwargs = recursive_kwargs or {}
-    return getattr(object_or_dict, method)(*r_args, **r_kwargs, **kwargs)
+    if isinstance(method, str):
+        return getattr(object_or_dict, method)(*r_args, **r_kwargs, **kwargs)
+    else:
+        return method(object_or_dict, *r_args, **r_kwargs, **kwargs)
 
 
 def is_exception(ex: Any) -> bool:

tests/_tests_dl_callbacks/config0.yml

@@ -38,6 +38,7 @@ stages:
         callback: CriterionCallback
       optimizer:
         callback: OptimizerCallback
+        use_fast_zero_grad: false
       accuracy:
         callback: AccuracyCallback
         accuracy_args: [1, 3, 5]