[hk] Expose hk.DepthwiseConv2D and change default padding to SAME

Also renames the home file to depthwise_conv. PiperOrigin-RevId: 294754337 Change-Id: Ic47e9f8bfac83095f3ddde8a08c42e5e6ac3fe4d
ultrons · Feb 12, 2020 · 62b19e6 · 62b19e6
1 parent 0092e1f
commit 62b19e6
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 29 deletions.
diff --git a/haiku/BUILD b/haiku/BUILD
@@ -28,6 +28,7 @@ py_library(
         "//haiku/_src:bias",
         "//haiku/_src:conv",
         "//haiku/_src:data_structures",
+        "//haiku/_src:depthwise_conv",
         "//haiku/_src:embed",
         "//haiku/_src:initializers",
         "//haiku/_src:layer_norm",

diff --git a/haiku/__init__.py b/haiku/__init__.py
@@ -47,6 +47,7 @@
 from haiku._src.conv import Conv2DTranspose
 from haiku._src.conv import Conv3D
 from haiku._src.conv import Conv3DTranspose
+from haiku._src.depthwise_conv import DepthwiseConv2D
 from haiku._src.embed import Embed
 from haiku._src.embed import EmbedLookupStyle
 from haiku._src.layer_norm import InstanceNorm
@@ -92,6 +93,7 @@
     "Conv3D",
     "Conv3DTranspose",
     "DeepRNN",
+    "DepthwiseConv2D",
     "EMAParamsTree",
     "Embed",
     "EmbedLookupStyle",

diff --git a/haiku/_src/BUILD b/haiku/_src/BUILD
@@ -18,7 +18,6 @@ hk_py_library(
         ":analytics",
         ":data_structures",
         ":typing",
-        ":utils",
         # pip: jax
         # pip: numpy
     ],
@@ -44,7 +43,6 @@ hk_py_library(
         ":module",
         ":utils",
         # pip: jax
-        # pip: numpy
     ],
 )
 
@@ -67,10 +65,8 @@ hk_py_library(
     srcs = ["module.py"],
     deps = [
         ":base",
-        ":data_structures",
         ":utils",
         # pip: jax
-        # pip: numpy
     ],
 )
 
@@ -139,25 +135,24 @@ hk_py_test(
 )
 
 hk_py_library(
-    name = "depthwiseconv",
-    srcs = ["depthwiseconv.py"],
+    name = "depthwise_conv",
+    srcs = ["depthwise_conv.py"],
     deps = [
         ":base",
         ":initializers",
         ":module",
-        ":pad",
         ":utils",
         # pip: jax
         # pip: numpy
     ],
 )
 
 hk_py_test(
-    name = "depthwiseconv_test",
-    srcs = ["depthwiseconv_test.py"],
+    name = "depthwise_conv_test",
+    srcs = ["depthwise_conv_test.py"],
     deps = [
         ":base",
-        ":depthwiseconv",
+        ":depthwise_conv",
         ":initializers",
         # pip: absl/testing:absltest
         # pip: absl/testing:parameterized
@@ -284,9 +279,6 @@ hk_py_test(
 hk_py_library(
     name = "utils",
     srcs = ["utils.py"],
-    deps = [
-        # pip: jax
-    ],
 )
 
 hk_py_test(
@@ -386,7 +378,6 @@ hk_py_library(
         ":initializers",
         ":module",
         # pip: jax
-        # pip: tree
     ],
 )
 
@@ -464,7 +455,6 @@ hk_py_library(
         ":basic",
         ":initializers",
         ":module",
-        # pip: enum
         # pip: jax
     ],
 )

diff --git a/haiku/_src/depthwiseconv.py → haiku/_src/depthwise_conv.py b/haiku/_src/depthwiseconv.py → haiku/_src/depthwise_conv.py
@@ -45,11 +45,11 @@ def __init__(self,
                channel_multiplier,
                kernel_shape,
                stride=1,
-               padding="VALID",
+               padding="SAME",
                with_bias=True,
                w_init=None,
                b_init=None,
-               data_format="channels_last",
+               data_format="NHWC",
                name=None):
     """Construct a 2D Depthwise Convolution.
 

diff --git a/haiku/_src/depthwiseconv_test.py → haiku/_src/depthwise_conv_test.py b/haiku/_src/depthwiseconv_test.py → haiku/_src/depthwise_conv_test.py
@@ -13,12 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for haiku._src.depthwiseconv."""
+"""Tests for haiku._src.depthwise_conv."""
 
 from absl.testing import absltest
 from absl.testing import parameterized
 from haiku._src import base
-from haiku._src import depthwiseconv
+from haiku._src import depthwise_conv
 from haiku._src import initializers
 from jax import random
 import jax.numpy as jnp
@@ -44,11 +44,11 @@ def f():
       data[0, :, :, 1] += 1
       data[0, :, :, 2] += 2
       data = jnp.array(data)
-      net = depthwiseconv.DepthwiseConv2D(
+      net = depthwise_conv.DepthwiseConv2D(
           channel_multiplier=1,
           kernel_shape=3,
           stride=1,
-          padding="Valid",
+          padding="VALID",
           with_bias=with_bias,
           data_format="channels_last",
           **create_constant_initializers(1.0, 1.0, with_bias))
@@ -76,11 +76,11 @@ def f():
       data[0, :, :, 1] += 1
       data[0, :, :, 2] += 2
       data = jnp.array(data)
-      net = depthwiseconv.DepthwiseConv2D(
+      net = depthwise_conv.DepthwiseConv2D(
           channel_multiplier=1,
           kernel_shape=3,
           stride=1,
-          padding="Same",
+          padding="SAME",
           with_bias=with_bias,
           data_format="channels_last",
           **create_constant_initializers(1.0, 0.0, with_bias))
@@ -97,11 +97,11 @@ def f():
       data[0, :, :, 1] += 1
       data[0, :, :, 2] += 2
       data = jnp.array(data)
-      net = depthwiseconv.DepthwiseConv2D(
+      net = depthwise_conv.DepthwiseConv2D(
           channel_multiplier=3,
           kernel_shape=3,
           stride=1,
-          padding="Valid",
+          padding="VALID",
           with_bias=with_bias,
           data_format="channels_last",
           **create_constant_initializers(1.0, 0.0, with_bias))
@@ -118,11 +118,11 @@ def f():
       data[0, 1, :, :] += 1
       data[0, 2, :, :] += 2
       data = jnp.array(data)
-      net = depthwiseconv.DepthwiseConv2D(
+      net = depthwise_conv.DepthwiseConv2D(
           channel_multiplier=1,
           kernel_shape=3,
           stride=1,
-          padding="Valid",
+          padding="VALID",
           with_bias=with_bias,
           data_format="channels_first",
           **create_constant_initializers(1.0, 1.0, with_bias))
@@ -147,11 +147,11 @@ def f():
       data[0, 1, :, :] += 1
       data[0, 2, :, :] += 2
       data = jnp.array(data)
-      net = depthwiseconv.DepthwiseConv2D(
+      net = depthwise_conv.DepthwiseConv2D(
           channel_multiplier=9,
           kernel_shape=3,
           stride=1,
-          padding="Valid",
+          padding="VALID",
           with_bias=with_bias,
           data_format="channels_first",
           **create_constant_initializers(1.0, 0.0, with_bias))

diff --git a/haiku/_src/integration/descriptors.py b/haiku/_src/integration/descriptors.py
@@ -149,6 +149,10 @@ class ModuleDescriptor(NamedTuple):
         name="Conv3DTranspose",
         create=lambda: hk.Conv3DTranspose(3, 3),
         shape=(BATCH_SIZE, 2, 2, 2, 2)),
+    ModuleDescriptor(
+        name="DepthwiseConv2D",
+        create=lambda: hk.DepthwiseConv2D(1, 3),
+        shape=(BATCH_SIZE, 2, 2, 2)),
 )