quant docs: add and clean up GroupNorm

Cleans up the quantized GroupNorm docstring and adds it to quantization docs. Test plan: * build on Mac OS and inspect ghstack-source-id: 4347125fe86d10d4f59f093dacc7b32cb1d2c6ea Pull Request resolved: #40343
pytorch · Jun 20, 2020 · ad62f2e · ad62f2e
1 parent 58143d9
commit ad62f2e
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 3 deletions.
diff --git a/docs/source/quantization.rst b/docs/source/quantization.rst
@@ -242,6 +242,7 @@ Layers for the quantization-aware training
 * :class:`~torch.nn.qat.Conv2d` — 2D convolution
 * :class:`~torch.nn.qat.Hardswish` — Hardswish
 * :class:`~torch.nn.qat.LayerNorm` — LayerNorm
+* :class:`~torch.nn.qat.GroupNorm` — GroupNorm
 
 ``torch.quantization``
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -354,6 +355,7 @@ Quantized version of standard NN layers.
   quantized representation of 6
 * :class:`~torch.nn.quantized.Hardswish` — Hardswish
 * :class:`~torch.nn.quantized.LayerNorm` — LayerNorm. *Note: performance on ARM is not optimized*.
+* :class:`~torch.nn.quantized.GroupNorm` — GroupNorm. *Note: performance on ARM is not optimized*.
 
 ``torch.nn.quantized.dynamic``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -714,6 +716,11 @@ LayerNorm
 .. autoclass:: LayerNorm
     :members:
 
+GroupNorm
+~~~~~~~~~~~~~~~
+.. autoclass:: GroupNorm
+    :members:
+
 
 torch.nn.quantized
 ----------------------------
@@ -802,6 +809,11 @@ LayerNorm
 .. autoclass:: LayerNorm
     :members:
 
+GroupNorm
+~~~~~~~~~~~~~~~
+.. autoclass:: GroupNorm
+    :members:
+
 torch.nn.quantized.dynamic
 ----------------------------
 

diff --git a/torch/nn/qat/modules/normalization.py b/torch/nn/qat/modules/normalization.py
@@ -3,8 +3,8 @@
 
 class GroupNorm(nn.GroupNorm):
     r"""
-    A GroupNorm module attached with FakeQuantize modules for both output
-    activation and weight, used for quantization aware training.
+    A GroupNorm module attached with FakeQuantize modules for output
+    activation, used for quantization aware training.
 
     Similar to `torch.nn.GroupNorm`, with FakeQuantize modules initialized to
     default.

diff --git a/torch/nn/quantized/modules/normalization.py b/torch/nn/quantized/modules/normalization.py
@@ -42,7 +42,12 @@ def from_float(cls, mod):
         return new_mod
 
 class GroupNorm(torch.nn.GroupNorm):
-    r"""This is the quantized version of `torch.nn.GroupNorm`.
+    r"""This is the quantized version of :class:`~torch.nn.GroupNorm`.
+
+    Additional args:
+        * **scale** - quantization scale of the output, type: double.
+        * **zero_point** - quantization zero point of the output, type: long.
+
     """
     __constants__ = ['num_groups', 'num_channels', 'eps', 'affine']