Add jax implementation of extract_patches.

PiperOrigin-RevId: 387610043

perceiver/io_processors.py

@@ -71,6 +71,72 @@ def space_to_depth(
         ' or rank 5 (batch, time, height, width, channels)')
 
 
+def extract_patches(images: jnp.ndarray,
+                    sizes: Sequence[int],
+                    strides: Sequence[int],
+                    rates: Sequence[int],
+                    padding: str = 'VALID') -> jnp.ndarray:
+  """Extract patches from images.
+
+  This function is a wrapper for jax.lax.conv_general_dilated_patches
+  to conforms to the same interface as tf.image.extract_patches.
+  The function extracts patches of shape sizes from the input images in the same
+  manner as a convolution with kernel of shape sizes, stride equal to strides,
+  and the given padding scheme.
+  The patches are stacked in the channel dimension.
+
+  Args:
+    images: input batch of images of shape [B, H, W, C].
+    sizes: size of extracted patches. Must be [1, size_rows, size_cols, 1].
+    strides: strides, must be [1, stride_rows, stride_cols, 1].
+    rates: sampling rate (as in dilated convolutions),
+      must be [1, rate_rows, rate_cols, 1].
+    padding: padding algorithm to use.
+  Returns:
+    Tensor of shape [B, patch_rows, patch_cols, size_rows * size_cols * C]
+  """
+
+  if len(sizes) != 4 or sizes[0] != 1 or sizes[3] != 1:
+    raise ValueError(
+        f'Shape of sizes must be [1, size_rows, size_cols, 1], got {sizes}.')
+  if len(strides) != 4 or strides[0] != 1 or strides[3] != 1:
+    raise ValueError(
+        f'Shape of strides must be [1, size_rows, size_cols, 1], '
+        f'got {strides}.')
+  if len(rates) != 4 or rates[0] != 1 or rates[3] != 1:
+    raise ValueError(
+        f'Shape of rates must be [1, size_rows, size_cols, 1], got {rates}.')
+  if images.ndim != 4:
+    raise ValueError(
+        f'Rank of images must be 4 (got tensor of shape {jnp.shape(images)})')
+  # Rearrange axes of images to NCHW for conv_general_dilated_patches
+  images = einops.rearrange(images, 'n h w c -> n c h w')
+  channels = images.shape[1]
+  patches = jax.lax.conv_general_dilated_patches(
+      images, sizes[1:-1], strides[1:-1], padding, rhs_dilation=rates[1:-1])
+  # conv_general_dilated_patches returns patches in channel-major order.
+  # Rearrange to match interface of tf.image.extract_patches.
+  patches = einops.rearrange(patches, 'n (c ph pw) h w -> n h w (ph pw c)',
+                             c=channels, ph=sizes[1], pw=sizes[2])
+  return patches
+
+
+def patches_for_flow(inputs: jnp.ndarray) -> jnp.ndarray:
+  """Extract 3x3x2 image patches for flow inputs."""
+
+  def pad_and_extract_patches(inputs):
+    padded_inputs = jnp.pad(inputs, [[0, 0], [1, 1], [1, 1], [0, 0]],
+                            mode='constant')
+    return extract_patches(
+        padded_inputs,
+        sizes=[1, 3, 3, 1],
+        strides=[1, 1, 1, 1],
+        padding='VALID',
+        rates=[1, 1, 1, 1])
+
+  return jax.vmap(pad_and_extract_patches, in_axes=1, out_axes=1)(inputs)
+
+
 #  ------------------------------------------------------------
 #  -------------------  Up/down-sampling  ---------------------
 #  ------------------------------------------------------------

perceiver/io_processors_test.py

@@ -0,0 +1,73 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for io_processors."""
+
+import numpy as np
+import tensorflow as tf
+
+from perceiver import io_processors
+
+
+def _create_test_image(shape):
+  image = np.arange(np.prod(np.array(shape)))
+  return np.reshape(image, shape)
+
+
+def test_space_to_depth_image():
+  image_shape = (2, 3 * 5, 3 * 7, 11)
+  image = _create_test_image(image_shape)
+  output = io_processors.space_to_depth(image, spatial_block_size=3)
+  assert output.shape == (2, 5, 7, 3 * 3 * 11)
+
+
+def test_space_to_depth_video():
+  image_shape = (2, 5 * 7, 3 * 11, 3 * 13, 17)
+  image = _create_test_image(image_shape)
+  output = io_processors.space_to_depth(image, spatial_block_size=3,
+                                        temporal_block_size=5)
+  assert output.shape == (2, 7, 11, 13, 5 * 3 * 3 * 17)
+
+
+def test_reverse_space_to_depth_image():
+  image_shape = (2, 5, 7, 3 * 3 * 11)
+  image = _create_test_image(image_shape)
+  output = io_processors.reverse_space_to_depth(image, spatial_block_size=3)
+  assert output.shape == (2, 3 * 5, 3 * 7, 11)
+
+
+def test_reverse_space_to_depth_video():
+  image_shape = (2, 7, 11, 13, 5 * 3 * 3 * 17)
+  image = _create_test_image(image_shape)
+  output = io_processors.reverse_space_to_depth(
+      image, spatial_block_size=3, temporal_block_size=5)
+  assert output.shape == (2, 5 * 7, 3 * 11, 3 * 13, 17)
+
+
+def test_extract_patches():
+  image_shape = (2, 5, 7, 3)
+  image = _create_test_image(image_shape)
+
+  sizes = [1, 2, 3, 1]
+  strides = [1, 1, 2, 1]
+  rates = [1, 2, 1, 1]
+
+  for padding in ["VALID", "SAME"]:
+    jax_patches = io_processors.extract_patches(
+        image, sizes=sizes, strides=strides, rates=rates, padding=padding)
+    tf_patches = tf.image.extract_patches(
+        image, sizes=sizes, strides=strides, rates=rates, padding=padding)
+    assert np.array_equal(
+        np.array(jax_patches),
+        tf_patches.numpy())