Code Update with meaningful docstrings

pytorchvideo/accelerator/deployment/common/model_transmuter.py

@@ -22,15 +22,27 @@ def _find_equivalent_efficient_module(
     module_name: str = "",
 ):
     """
+    Searches for an equivalent efficientBlock that can replace the given `module_input`
+    within the efficient_block_transmuter_list.
+
     Given module_input, search through efficient_block_registry to see whether the
     module_input can be replaced with equivalent efficientBlock. Returns None if no
     equivalent efficientBlock is found, else returns an instance of equivalent
     efficientBlock.
+
     Args:
-        module_input (nn.Module): module to be replaced by equivalent efficientBlock
-        efficient_block_transmuter_list (list): a transmuter list that contains transmuter
-            functions for available efficientBlocks
-        module_name (str): name of module_input in original model
+        module_input (nn.Module): The module to be replaced by an equivalent efficientBlock.
+        efficient_block_transmuter_list (list): A list containing transmuter functions for
+            available efficientBlocks.
+        module_name (str): The name of `module_input` in the original model.
+
+    Returns:
+        nn.Module or None: An instance of the equivalent efficientBlock if found; otherwise, None.
+
+    This function iterates through the `efficient_block_transmuter_list` and applies each transmuter
+    function to `module_input`. If an equivalent efficientBlock is found, it is added to the
+    `eq_module_hit_list`. If multiple matches are found, a warning is logged, and the one with
+    the highest priority is chosen. If no matches are found, None is returned.
     """
     eq_module_hit_list = []
     for iter_func in efficient_block_transmuter_list:
@@ -56,13 +68,20 @@ def transmute_model(
     prefix: str = "",
 ):
     """
-    Recursively goes through user input model and replace module in place with available
-    equivalent efficientBlock for target device.
+    Recursively goes through the user input model and replaces modules in place with
+    equivalent efficientBlocks suitable for the target device.
+
     Args:
-        model (nn.Module): user input model to be transmuted
-        target_device (str): name of target device, used to access transmuter list in
-            EFFICIENT_BLOCK_TRANSMUTER_REGISTRY
-        prefix (str): name of current hierarchy in user model
+        model (nn.Module): The user input model to be transmuted.
+        target_device (str): The name of the target device, used to access the transmuter
+            list in EFFICIENT_BLOCK_TRANSMUTER_REGISTRY.
+        prefix (str): The name of the current hierarchy in the user model.
+
+    This function recursively traverses the input `model`, examining each child module.
+    It attempts to find an equivalent efficientBlock for each module and replaces it
+    in the model if an equivalent is found. The replacement is logged for reference.
+
+    Note: Make sure the target device is registered in the EFFICIENT_BLOCK_TRANSMUTER_REGISTRY.
     """
     assert (
         target_device in EFFICIENT_BLOCK_TRANSMUTER_REGISTRY

pytorchvideo/accelerator/deployment/mobile_cpu/transmuter/transmuter_mobile_cpu.py

@@ -12,11 +12,21 @@
 
 def transmute_Conv3dPwBnAct(input_module: nn.Module):
     """
-    Given an input_module, transmutes it into a equivalent Conv3dPwBnAct. Returns None
-    if no equivalent Conv3dPwBnAct is found, else returns an instance of equivalent
-    Conv3dPwBnAct.
+    Transmutes the given `input_module` into an equivalent Conv3dPwBnAct module if applicable.
+
     Args:
-        input_module (nn.Module): input module to find an equivalent Conv3dPwBnAct
+        input_module (nn.Module): The input module to find an equivalent Conv3dPwBnAct for.
+
+    Returns:
+        Conv3dPwBnAct or None: An instance of the equivalent Conv3dPwBnAct module if found;
+        otherwise, None.
+
+    This function checks if `input_module` is an instance of nn.Conv3d and if it matches specific
+    criteria, such as kernel size, groups, stride, padding, and dilation. If the criteria are met,
+    it creates and returns an equivalent Conv3dPwBnAct module, copying the weights if necessary.
+
+    Note: Conv3dPwBnAct is a module that combines a 3D pointwise convolution with batch normalization
+    and activation functions.
     """
     if not isinstance(input_module, nn.Conv3d):
         return None
@@ -42,11 +52,22 @@ def transmute_Conv3dPwBnAct(input_module: nn.Module):
 
 def transmute_Conv3d3x3x3DwBnAct(input_module: nn.Module):
     """
-    Given an input_module, transmutes it into a equivalent Conv3d3x3x3DwBnAct. Returns
-    None if no equivalent Conv3d3x3x3DwBnAct is found, else returns an instance of
-    equivalent Conv3d3x3x3DwBnAct.
+    Transmutes the given `input_module` into an equivalent Conv3d3x3x3DwBnAct module if applicable.
+
     Args:
-        input_module (nn.Module): input module to find an equivalent Conv3d3x3x3DwBnAct
+        input_module (nn.Module): The input module to find an equivalent Conv3d3x3x3DwBnAct for.
+
+    Returns:
+        Conv3d3x3x3DwBnAct or None: An instance of the equivalent Conv3d3x3x3DwBnAct module if found;
+        otherwise, None.
+
+    This function checks if `input_module` is an instance of nn.Conv3d and if it matches specific
+    criteria, such as kernel size, in_channels, groups, stride, padding, padding_mode, and dilation.
+    If the criteria are met, it creates and returns an equivalent Conv3d3x3x3DwBnAct module, copying
+    the weights if necessary.
+
+    Note: Conv3d3x3x3DwBnAct is a module that combines a 3D 3x3x3 depthwise convolution with batch
+    normalization and activation functions.
     """
     if not isinstance(input_module, nn.Conv3d):
         return None

pytorchvideo/accelerator/deployment/mobile_cpu/utils/model_conversion.py

@@ -17,14 +17,21 @@ def _add_input_tensor_size_lut_hook(
     base_name: str = "",
 ) -> None:
     """
-    This helper function recursively goes through all modules in a network, registers
-    forward hook function to each module. The hook function records the input tensor
-    size in forward in input_tensor_size_lut[base_name].
+    Recursively adds a forward hook to each module in a network, recording input tensor sizes
+    in the provided input_tensor_size_lut dictionary.
+
     Args:
-        module (nn.Module): input module to add hook recursively.
-        input_tensor_size_lut (dict): lut to record input tensor size for hook function.
-        hook_handle_list (list): a list to contain hook handles.
-        base_name (str): name for module input.
+        module (nn.Module): The input module to add hooks to, recursively.
+        input_tensor_size_lut (dict): A dictionary to record input tensor sizes for the hook function.
+        hook_handle_list (list): A list to contain hook handles.
+        base_name (str): The base name for the input module.
+
+    This helper function iterates through the input `module` and its children, registering a forward
+    hook for each module. The hook function records the input tensor size for the module in the
+    `input_tensor_size_lut` dictionary using the `base_name` as the key.
+
+    Note: Forward hooks are useful for monitoring and analyzing the input tensor sizes as they pass
+    through each module in a neural network.
     """
 
     def hook_fn(_, _in, _out):
@@ -51,21 +58,27 @@ def _convert_module(
     native_conv3d_op_qnnpack: bool = False,
 ) -> None:
     """
-    This helper function recursively goes through sub-modules in a network. If current
-    module is a efficient block (instance of EfficientBlockBase) with convert() method,
-    its convert() method will be called, and the input tensor size (needed by efficient
-    blocks for mobile cpu) will be provided by matching module name in
-    input_tensor_size_lut.
-    Otherwise if the input module is a non efficient block, this function will try to go
-    through child modules of input module to look for any efficient block in lower
-    hierarchy.
+    Recursively traverses sub-modules in a neural network and performs module conversion
+    if applicable. For efficient blocks (instances of EfficientBlockBase) with a 'convert'
+    method, it calls the 'convert' method with the input tensor size obtained from the
+    'input_tensor_size_lut'. If the module is not an efficient block, it explores child
+    modules to find efficient blocks in lower hierarchy.
+
     Args:
-        module (nn.Module): input module for convert.
-        input_tensor_size_lut (dict): input tensor size look-up table.
-        base_name (str): module name for input module.
-        convert_for_quantize (bool): whether this module is intended to be quantized.
-        native_conv3d_op_qnnpack (bool): whether the QNNPACK version has native int8
-            Conv3d.
+        module (nn.Module): The input module for conversion.
+        input_tensor_size_lut (dict): A dictionary containing input tensor sizes for reference.
+        base_name (str): The name of the module.
+        convert_for_quantize (bool): Whether this module is intended for quantization.
+        native_conv3d_op_qnnpack (bool): Whether the QNNPACK version has native int8 Conv3d.
+
+    This helper function is designed for recursively exploring a neural network and converting
+    specific modules, such as efficient blocks. If a module is an instance of EfficientBlockBase
+    and has a 'convert' method, it calls the 'convert' method with the input tensor size from
+    the 'input_tensor_size_lut'. If the module is not an efficient block, it continues to explore
+    its child modules in search of efficient blocks in lower hierarchies.
+
+    Note: Module conversion is a common step in optimizing and adapting neural networks for
+    specific hardware or use cases, such as mobile CPUs.
     """
     if isinstance(module, EfficientBlockBase):
         module.convert(
@@ -91,19 +104,31 @@ def convert_to_deployable_form(
     native_conv3d_op_qnnpack: bool = False,
 ) -> nn.Module:
     """
-    This function takes an input model, and returns a deployable model copy.
+    Converts an input model into a deployable form and returns a copy of the modified model.
+
     Args:
-        model (nn.Module): input model for conversion. The model can include a mix of
-            efficient blocks (instances of EfficientBlockBase) and non efficient blocks.
-            The efficient blocks will be converted by calling its convert() method, while
-            other blocks will stay unchanged.
-        input_tensor (torch.Tensor): input tensor for model. Note current conversion for
-            deployable form in mobile cpu only works for single input tensor size (i.e.,
-            the future input tensor to converted model should have the same size as
-            input_tensor specified here).
-        convert_for_quantize (bool): whether this module is intended to be quantized.
-        native_conv3d_op_qnnpack (bool): whether the QNNPACK version has native int8
-            Conv3d.
+        model (nn.Module): The input model for conversion. The model can consist of a mix
+            of efficient blocks (instances of EfficientBlockBase) and non-efficient blocks.
+            Efficient blocks are converted using their `convert()` method, while other
+            blocks remain unchanged.
+        input_tensor (torch.Tensor): The input tensor used for the model. The conversion for
+            deployable form on mobile CPU is designed for a single input tensor size. The
+            future input tensor to the converted model should match the size of the
+            `input_tensor` specified here.
+        convert_for_quantize (bool): Indicates whether this module is intended for quantization.
+        native_conv3d_op_qnnpack (bool): Specifies whether the QNNPACK version has native
+            int8 Conv3d support.
+
+    Returns:
+        nn.Module: A copy of the input model converted into a deployable form.
+
+    This function prepares the input model for deployment by performing the following steps:
+    1. Captures input tensor sizes during forward pass.
+    2. Executes a forward pass to record input tensor sizes.
+    3. Removes forward hooks used for input tensor size capture.
+    4. Creates a deep copy of the input model for conversion.
+    5. Converts the copied model by applying the `_convert_module` function.
+    6. Returns the converted model suitable for deployment.
     """
     input_tensor_size_lut = {}
     hook_handle_list = []

pytorchvideo/accelerator/efficient_blocks/efficient_block_base.py

@@ -7,29 +7,37 @@
 
 class EfficientBlockBase(nn.Module):
     """
-    PyTorchVideo/accelerator provides a set of efficient blocks
-    that have optimal efficiency for each target hardware device.
+    The EfficientBlockBase is the foundation for efficient blocks provided by PyTorchVideo's accelerator.
+    These efficient blocks are designed for optimal efficiency on various target hardware devices.
 
     Each efficient block has two forms:
-    - original form: this form is for training. When efficient block is instantiated,
-        it is in this original form.
-    - deployable form: this form is for deployment. Once the network is ready for
-        deploy, it can be converted into deployable form for efficient execution
-        on target hardware. One block is transformed into deployable form by calling
-        convert() method. By conversion to deployable form,
-        various optimization (operator fuse, kernel optimization, etc.) are applied.
-
-    EfficientBlockBase is the base class for efficient blocks.
-    All efficient blocks should inherit this base class
-    and implement following methods:
-    - forward(): same as required by nn.Module
-    - convert(): called to convert block into deployable form
+    - Original Form: This form is used during training. When an efficient block is instantiated,
+        it is in its original form.
+    - Deployable Form: This form is for deployment. Once the network is prepared for deployment,
+        it can be converted into the deployable form to enable efficient execution on the target hardware.
+        Conversion to the deployable form involves various optimizations such as operator fusion
+        and kernel optimization.
+
+    All efficient blocks must inherit from this base class and implement the following methods:
+    - forward(): This method serves the same purpose as required by nn.Module.
+    - convert(): Called to transform the block into its deployable form.
+
+    Subclasses of EfficientBlockBase should provide implementations for these methods to tailor
+    the behavior of the efficient block for specific use cases and target hardware.
+
+    Note: This class is abstract, and its methods must be implemented in derived classes.
     """
 
     @abstractmethod
     def convert(self):
+        """
+        Abstract method to convert the efficient block into its deployable form.
+        """
         pass
 
     @abstractmethod
     def forward(self):
+        """
+        Abstract method for the forward pass of the efficient block.
+        """
         pass

pytorchvideo/accelerator/efficient_blocks/no_op_convert_block.py

@@ -7,14 +7,20 @@
 
 class NoOpConvertBlock(EfficientBlockBase):
     """
-    This class provides an interface with EfficientBlockBase for modules that do not
-    need convert.
+    A class that provides an interface with EfficientBlockBase for modules that do not
+    require conversion.
+
     Args:
-        model (nn.Module): NoOpConvertBlock takes model as input and generate a wrapper
-            instance of EfficientBlockBase with same functionality as model, with no change
-            applied when convert() is called.
-    """
+        model (nn.Module): NoOpConvertBlock takes a model as input and generates a wrapper
+            instance of EfficientBlockBase with the same functionality as the model. When
+            `convert()` is called on this instance, no changes are applied.
 
+    This class is designed for modules that do not need any conversion when integrated into
+    an EfficientBlockBase. It takes an existing `model` and acts as a pass-through, forwarding
+    input directly to the underlying model during the `forward` pass. When `convert()` is
+    called, it simply does nothing, ensuring that no modifications are made to the model.
+    """
+
     def __init__(self, model: nn.Module):
         super().__init__()
         self.model = model