From 595d0727b2537ffcaeb97f636f2b8aadc8e45c04 Mon Sep 17 00:00:00 2001 From: David Ding Date: Thu, 29 Jul 2021 18:05:07 +0100 Subject: [PATCH] Add jax implementation of extract_patches. PiperOrigin-RevId: 387610043 --- perceiver/io_processors.py | 66 +++++++++++++++++++++++++++++ perceiver/io_processors_test.py | 73 +++++++++++++++++++++++++++++++++ 2 files changed, 139 insertions(+) create mode 100644 perceiver/io_processors_test.py diff --git a/perceiver/io_processors.py b/perceiver/io_processors.py index 38955e6..9696656 100644 --- a/perceiver/io_processors.py +++ b/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 --------------------- # ------------------------------------------------------------ diff --git a/perceiver/io_processors_test.py b/perceiver/io_processors_test.py new file mode 100644 index 0000000..3167ed1 --- /dev/null +++ b/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())