docs(mge): modify some docstring about conv and tools

GitOrigin-RevId: 4833e78d71

imperative/python/megengine/functional/nn.py

@@ -513,7 +513,22 @@ def local_conv2d(
     dilation: Union[int, Tuple[int, int]] = 1,
     conv_mode="cross_correlation",
 ):
-    r"""Applies spatial 2D convolution over an groupped channeled image with untied kernels."""
+    r"""Applies a spatial convolution with untied kernels over an groupped channeled input 4D tensor.
+    It is also known as the locally connected layer.
+
+    Args:
+        inp: input feature map.
+        weight: convolution kernel.
+            weight usually has shape ``(out_channels, in_channels, height, width)``.
+        bias: bias added to the result of convolution (if given).
+        stride: stride of the 2D convolution operation. Default: 1
+        padding: size of the paddings added to the input on both sides of its
+            spatial dimensions. Only zero-padding is supported. Default: 0
+        dilation: dilation of the 2D convolution operation. Default: 1
+
+    Returns:
+        output tensor.
+    """
     assert (
         conv_mode.lower() == "cross_correlation"
         or conv_mode.name == "CROSS_CORRELATION"
@@ -529,6 +544,7 @@ def local_conv2d(
     if weight.dtype != dtype:
         weight = weight.astype(dtype)
 
+    # local conv only support "dense" mode, but weight could contain group dimension.
     op = builtin.GroupLocal(
         stride_h=stride_h,
         stride_w=stride_w,
@@ -568,6 +584,11 @@ def conv_transpose3d(
         padding: size of the paddings added to the input on all sides of its
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 3D convolution operation. Default: 1
+        groups: number of groups into which the input and output channels are divided,
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
+            ``in_channels`` and ``out_channels`` must be divisible by groups,
+            and the shape of weight should be ``(groups, in_channels // groups,
+            out_channels // groups, depth, height, width)``. Default: 1
 
     Returns:
         output tensor.

imperative/python/megengine/module/conv.py

@@ -128,11 +128,11 @@ class Conv1d(_ConvNd):
         padding: size of the paddings added to the input on both sides of its
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 1D convolution operation. Default: 1
-        groups: number of groups into which the input and output channels are divided,
+        groups: number of groups to divide input and output channels into,
             so as to perform a "grouped convolution". When ``groups`` is not 1,
             ``in_channels`` and ``out_channels`` must be divisible by ``groups``,
-            and there would be an extra dimension at the beginning of the weight's
-            shape. Default: 1
+            and the shape of weight should be ``(groups, out_channel // groups,
+            in_channels // groups, kernel_size)``. Default: 1
         bias: whether to add a bias onto the result of convolution. Default: True
         conv_mode: Supports `cross_correlation`. Default: `cross_correlation`
         compute_mode: When set to "default", no special requirements will be
@@ -290,10 +290,10 @@ class Conv2d(_ConvNd):
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 2D convolution operation. Default: 1
         groups: number of groups into which the input and output channels are divided,
-            so as to perform a "grouped convolution". When ``groups`` is not 1,
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
             ``in_channels`` and ``out_channels`` must be divisible by ``groups``,
-            and there would be an extra dimension at the beginning of the weight's
-            shape. Default: 1
+            and the shape of weight should be ``(groups, out_channel // groups,
+            in_channels // groups, height, width)``. Default: 1
         bias: whether to add a bias onto the result of convolution. Default: True
         conv_mode: Supports `cross_correlation`. Default: `cross_correlation`
         compute_mode: When set to "default", no special requirements will be
@@ -436,10 +436,10 @@ class Conv3d(_ConvNd):
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 3D convolution operation. Default: 1
         groups: number of groups into which the input and output channels are divided,
-            so as to perform a "grouped convolution". When ``groups`` is not 1,
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
             ``in_channels`` and ``out_channels`` must be divisible by ``groups``,
-            and there would be an extra dimension at the beginning of the weight's
-            shape. Default: 1
+            and the shape of weight should be ``(groups, out_channel // groups,
+            in_channels // groups, depth, height, width)``. Default: 1
         bias: whether to add a bias onto the result of convolution. Default: True
         conv_mode: Supports `cross_correlation`. Default: `cross_correlation`
 
@@ -560,10 +560,10 @@ class ConvTranspose2d(_ConvNd):
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 2D convolution operation. Default: 1
         groups: number of groups into which the input and output channels are divided,
-            so as to perform a "grouped convolution". When ``groups`` is not 1,
-            ``in_channels`` and ``out_channels`` must be divisible by ``groups``,
-            and there would be an extra dimension at the beginning of the weight's
-            shape. Default: 1
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
+            ``in_channels`` and ``out_channels`` must be divisible by groups,
+            and the shape of weight should be ``(groups, in_channels // groups,
+            out_channels // groups, height, width)``. Default: 1
         bias: wether to add a bias onto the result of convolution. Default: True
             conv_mode: Supports `cross_correlation`. Default: `cross_correlation`
         compute_mode: When set to "default", no special requirements will be
@@ -667,6 +667,7 @@ class LocalConv2d(Conv2d):
         stride: stride of the 2D convolution operation. Default: 1
         padding: size of the paddings added to the input on both sides of its
             spatial dimensions. Only zero-padding is supported. Default: 0
+        dilation: dilation of the 2D convolution operation. Default: 1
         groups: number of groups into which the input and output channels are divided,
             so as to perform a "grouped convolution". When ``groups`` is not 1,
             ``in_channels`` and ``out_channels`` must be divisible by ``groups``. Default: 1
@@ -759,10 +760,10 @@ class DeformableConv2d(_ConvNd):
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 2D convolution operation. Default: 1
         groups: number of groups into which the input and output channels are divided,
-            so as to perform a "grouped convolution". When ``groups`` is not 1,
-            ``in_channels`` and ``out_channels`` must be divisible by ``groups``,
-            and there would be an extra dimension at the beginning of the weight's
-            shape. Default: 1
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
+            ``in_channels`` and ``out_channels`` must be divisible by groups,
+            and the shape of weight should be ``(groups, out_channel // groups,
+            in_channels // groups, height, width)``. Default: 1
         bias: whether to add a bias onto the result of convolution. Default: True
         conv_mode: Supports `cross_correlation`. Default: `cross_correlation`
         compute_mode: When set to "default", no special requirements will be
@@ -875,6 +876,11 @@ class ConvTranspose3d(_ConvNd):
         padding: size of the paddings added to the input on all sides of its
             spatial dimensions. Only zero-padding is supported. Default: 0
         dilation: dilation of the 3D convolution operation. Default: 1
+        groups: number of groups into which the input and output channels are divided,
+            so as to perform a ``grouped convolution``. When ``groups`` is not 1,
+            ``in_channels`` and ``out_channels`` must be divisible by groups,
+            and the shape of weight should be ``(groups, in_channels // groups,
+            out_channels // groups, depth, height, width)``. Default: 1
         bias: wether to add a bias onto the result of convolution. Default: True
 
     Note:
@@ -890,8 +896,8 @@ class ConvTranspose3d(_ConvNd):
         stride: Union[int, Tuple[int, int, int]] = 1,
         padding: Union[int, Tuple[int, int, int]] = 0,
         dilation: Union[int, Tuple[int, int, int]] = 1,
-        bias: bool = True,
         groups: int = 1,
+        bias: bool = True,
     ):
         kernel_size = _triple_nonzero(kernel_size)
         stride = _triple_nonzero(stride)

imperative/python/megengine/tools/network_visualize.py

@@ -53,8 +53,11 @@ def visualize(
     Args:
       model_path: dir path for megengine dumped model.
       log_path: dir path for tensorboard graph log.
-      input: user defined input data for running model and calculating stats, alternative with inp_dict, used when the model has only one input.
-      inp_dict: input dict for running model and calculating stats, alternative with input, used when the model has more than one input. When both input and inp_dict are None, a random input will be used.
+      input: user defined input data for running model and calculating stats,
+        alternative with inp_dict, used when the model has only one input.
+      inp_dict: input dict for running model and calculating stats, alternative with
+        input, used when the model has more than one input.
+        When both input and inp_dict are None, a random input will be used.
       cal_params: whether calculate and record params size.
       cal_flops: whether calculate and record op flops.
       cal_activations: whether calculate and record op activations.