Add ScaleSpaceFlow network

Author: chyomin06

.clang-format

@@ -23,7 +23,7 @@ AlwaysBreakBeforeMultilineStrings: false
 AlwaysBreakTemplateDeclarations: MultiLine
 BinPackArguments: true
 BinPackParameters: true
-BraceWrapping:   
+BraceWrapping:
   AfterCaseLabel:  false
   AfterClass:      false
   AfterControlStatement: false
@@ -60,12 +60,12 @@ DerivePointerAlignment: false
 DisableFormat:   false
 ExperimentalAutoDetectBinPacking: false
 FixNamespaceComments: true
-ForEachMacros:   
+ForEachMacros:
   - foreach
   - Q_FOREACH
   - BOOST_FOREACH
 IncludeBlocks:   Preserve
-IncludeCategories: 
+IncludeCategories:
   - Regex:           '^"(llvm|llvm-c|clang|clang-c)/'
     Priority:        2
   - Regex:           '^(<|"(gtest|gmock|isl|json)/)'
@@ -117,10 +117,9 @@ SpacesInCStyleCastParentheses: false
 SpacesInParentheses: false
 SpacesInSquareBrackets: false
 Standard:        Cpp11
-StatementMacros: 
+StatementMacros:
   - Q_UNUSED
   - QT_REQUIRE_VERSION
 TabWidth:        8
 UseTab:          Never
 ...
-

.flake8

@@ -11,7 +11,7 @@ per-file-ignores =
 
 max-line-length = 88
 
-# maximum McCabe complexity 
+# maximum McCabe complexity
 max-complexity = 12
 
 exclude =

.gitignore

@@ -96,11 +96,11 @@ ipython_config.py
 .python-version
 
 # pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in 
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in
 #   version control.
-#   However, in case of collaboration, if having platform-specific dependencies 
+#   However, in case of collaboration, if having platform-specific dependencies
 #   or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don’t 
+#   having no cross-platform support, pipenv may install dependencies that don’t
 #   work, or not
 #   install all needed dependencies.
 #Pipfile.lock

compressai/layers/layers.py

@@ -18,6 +18,7 @@
 import torch.nn as nn
 
 from torch import Tensor
+from torch.autograd import Function
 
 from .gdn import GDN
 
@@ -29,6 +30,7 @@
     "ResidualBlockWithStride",
     "conv3x3",
     "subpel_conv3x3",
+    "qrelu"
 ]
 
 
@@ -225,3 +227,50 @@ def forward(self, x: Tensor) -> Tensor:
         out = a * torch.sigmoid(b)
         out += identity
         return out
+
+class qrelu(Function):
+    """QReLU
+
+    Clamping input with given bit-depth range. 
+    Suppose that input data presents integer through the integer networks
+    otherwise fraction values of input bypass without rounding operation for floating point networks.
+
+    Pre-computed scale with gamma function is used for backward computation.
+    
+    More details can be found `"Integer networks for data compression with latent-variable models"
+    
+    #_______`_, Balle et al. in 2019
+
+    Args:
+        input : a tensor data
+        bit_depth : bit-depth for clamping input
+        beta : 
+    """
+
+    @staticmethod
+    def forward(ctx, input, bit_depth, beta):
+        ctx.alpha = 0.9943258522851727
+        ctx.beta = beta
+        ctx.max_value = 2 ** bit_depth - 1
+        ctx.save_for_backward(input)
+
+        return input.clamp(min=0, max=ctx.max_value)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = None
+        (input,) = ctx.saved_tensors
+
+        grad_input = grad_output.clone()
+        grad_sub = (
+            torch.exp(
+                (-ctx.alpha ** ctx.beta)
+                * torch.abs(2.0 * input / ctx.max_value - 1) ** ctx.beta
+            )
+            * grad_output.clone()
+        )
+
+        grad_input[input < 0] = grad_sub[input < 0]
+        grad_input[input > ctx.max_value] = grad_sub[input > ctx.max_value]
+
+        return grad_input, None, None

compressai/models/__init__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 from .google import *
-from .waseda import *
+from .waseda import *