Add Registry to batch transform callback

catalyst/callbacks/batch_transform.py

@@ -1,6 +1,7 @@
 from typing import Callable, List, Union
 
 from catalyst.core import Callback, CallbackOrder, IRunner
+from catalyst.registry import REGISTRY
 
 
 def _tuple_wrapper(transform: Callable):
@@ -17,7 +18,7 @@ class BatchTransformCallback(Callback):
     Preprocess your batch with specified function.
 
     Args:
-        transform (Callable): Function to apply.
+        transform (Callable, str): Function to apply. If string will get function from registry.
         scope (str): ``"on_batch_end"`` (post-processing model output) or
             ``"on_batch_start"`` (pre-processing model input).
         input_key (Union[List[str], str], optional): Keys in batch dict to apply function.
@@ -165,12 +166,22 @@ def __len__(self):
                     )
                 ],
             )
+        .. code-block:: yaml
+
+            ...
+            callbacks:
+                transform:
+                    _target_: BatchTransformCallback
+                    transform: catalyst.ToTensor
+                    scope: on_batch_start
+                    input_key: features
+
 
     """
 
     def __init__(
         self,
-        transform: Callable,
+        transform: Union[Callable, str],
         scope: str,
         input_key: Union[List[str], str] = None,
         output_key: Union[List[str], str] = None,
@@ -179,7 +190,7 @@ def __init__(
         Preprocess your batch with specified function.
 
         Args:
-            transform (Callable): Function to apply.
+            transform (Callable, str): Function to apply. If string will get function from registry.
             scope (str): ``"on_batch_end"`` (post-processing model output) or
                 ``"on_batch_start"`` (pre-processing model input).
             input_key (Union[List[str], str], optional): Keys in batch dict to apply function.
@@ -206,7 +217,9 @@ def __init__(
         output_key = output_key or input_key
         if output_key is not None:
             if input_key is None:
-                raise TypeError("You should define input_key in " "case if output_key is not None")
+                raise TypeError(
+                    "You should define input_key in " "case if output_key is not None"
+                )
             if not isinstance(output_key, (list, str)):
                 raise TypeError("output key should be str or a list of str.")
             if isinstance(output_key, str):
@@ -219,6 +232,8 @@ def __init__(
             raise TypeError('Expected scope to be on of the ["on_batch_end", "on_batch_start"]')
         self.input_key = input_key
         self.output_key = output_key
+        if isinstance(transform, str):
+            transform = REGISTRY.get(transform)
         self.transform = transform
 
     def _handle_value(self, runner):