diff --git a/oneflow/python/nn/modules/conv.py b/oneflow/python/nn/modules/conv.py
index 0d8398cd08009c0f43cc1530c5cf21ae6c8eb198..769912ee3631f2e28ed27b56cf492ad5f1da43e0 100644
--- a/oneflow/python/nn/modules/conv.py
+++ b/oneflow/python/nn/modules/conv.py
@@ -40,6 +40,119 @@ class ConvUtil(object):
@oneflow_export("nn.Conv2d")
@experimental_api
class Conv2d(Module):
+ r"""Applies a 2D convolution over an input signal composed of several input
+ planes.
+
+ In the simplest case, the output value of the layer with input size
+ :math:`(N, C_{\text{in}}, H, W)` and output :math:`(N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})`
+ can be precisely described as:
+
+ .. math::
+ \text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) +
+ \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)
+
+
+ where :math:`\star` is the valid 2D `cross-correlation`_ operator,
+ :math:`N` is a batch size, :math:`C` denotes a number of channels,
+ :math:`H` is a height of input planes in pixels, and :math:`W` is
+ width in pixels.
+
+
+ * :attr:`stride` controls the stride for the cross-correlation, a single
+ number or a tuple.
+
+ * :attr:`padding` controls the amount of implicit padding on both
+ sides for :attr:`padding` number of points for each dimension.
+
+ * :attr:`dilation` controls the spacing between the kernel points; also
+ known as the à trous algorithm. It is harder to describe, but this `link`_
+ has a nice visualization of what :attr:`dilation` does.
+
+ * :attr:`groups` controls the connections between inputs and outputs. :attr:`in_channels`
+ and :attr:`out_channels` must both be divisible by :attr:`groups`. For example,
+
+ * At groups=1, all inputs are convolved to all outputs.
+ * At groups=2, the operation becomes equivalent to having two conv
+ layers side by side, each seeing half the input channels
+ and producing half the output channels, and both subsequently
+ concatenated.
+ * At groups= :attr:`in_channels`, each input channel is convolved with
+ its own set of filters (of size
+ :math:`\frac{\text{out_channels}}{\text{in_channels}}`).,
+
+ The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:
+
+ - a single ``int`` -- in which case the same value is used for the height and width dimension
+ - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+ and the second `int` for the width dimension
+
+ Note:
+ When `groups == in_channels` and `out_channels == K * in_channels`,
+ where `K` is a positive integer, this operation is also known as a "depthwise convolution".
+
+ In other words, for an input of size :math:`(N, C_{in}, L_{in})`,
+ a depthwise convolution with a depthwise multiplier `K` can be performed with the arguments
+ :math:`(C_\text{in}=C_\text{in}, C_\text{out}=C_\text{in} \times \text{K}, ..., \text{groups}=C_\text{in})`.
+
+
+ Args:
+ in_channels (int): Number of channels in the input image
+ out_channels (int): Number of channels produced by the convolution
+ kernel_size (int or tuple): Size of the convolving kernel
+ stride (int or tuple, optional): Stride of the convolution. Default: 1
+ padding (int or tuple, optional): Zero-padding added to both sides of
+ the input. Default: 0
+ padding_mode (string, optional): ``'zeros'``, ``'reflect'``,
+ ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+ dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+ groups (int, optional): Number of blocked connections from input
+ channels to output channels. Default: 1
+ bias (bool, optional): If ``True``, adds a learnable bias to the
+ output. Default: ``True``
+
+ Shape:
+ - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
+ - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where
+
+ .. math::
+ H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0]
+ \times (\text{kernel_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
+
+ .. math::
+ W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1]
+ \times (\text{kernel_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
+
+ Attr:
+ - weight (Tensor): the learnable weights of the module of shape
+ :math:`(\text{out_channels}, \frac{\text{in_channels}}{\text{groups}},`
+ :math:`\text{kernel_size[0]}, \text{kernel_size[1]})`.
+ The values of these weights are sampled from
+ :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+ :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel_size}[i]}`
+
+ - bias (Tensor): the learnable bias of the module of shape
+ (out_channels). If :attr:`bias` is ``True``,
+ then the values of these weights are
+ sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+ :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel_size}[i]}`
+
+ For example:
+
+ .. code-block:: python
+
+ import oneflow as flow
+
+ m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
+ input = flow.randn(20, 16, 50, 100)
+ output = m(input)
+
+ .. _cross-correlation:
+ https://en.wikipedia.org/wiki/Cross-correlation
+
+ .. _link:
+ https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+ """
+
def __init__(
self,
in_channels: int,
diff --git a/oneflow/python/nn/modules/linear.py b/oneflow/python/nn/modules/linear.py
index ec73b801007c21b65cb4a44507b60614b53d1874..87f9b8611388fa54a58e2cdf736917ab57dfe2c2 100644
--- a/oneflow/python/nn/modules/linear.py
+++ b/oneflow/python/nn/modules/linear.py
@@ -76,10 +76,10 @@ class Linear(Module):
- Output: :math:`(N, *, H_{out})` where all but the last dimension
are the same shape as the input and :math:`H_{out} = {out\_features}`.
- Attributes:
- weight: the learnable weights of the module of shape :math:`({out\_features}, {in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`(k = 1 / {in\_features})`
+ Attr:
+ - :attr:`weight`: the learnable weights of the module of shape :math:`({out\_features}, {in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`(k = 1 / {in\_features})`
- bias: the learnable bias of the module of shape :math:`({out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`(k = 1 / {in\_features})`
+ - :attr:`bias`: the learnable bias of the module of shape :math:`({out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`(k = 1 / {in\_features})`
For example: