Add `affine_transform` op to all backends #477

james77777778 · 2023-07-14T09:13:43Z

Related to #460

This PR introduces the affine operation in 3 backends. This op can be useful when implementing KPLs using pure keras-core, such as RandomRotation, RandomTranslation, etc.

It is worth noting that the implementations of F.affine_grid and F.grid_sample in the torch backend differ significantly from TF and JAX. As a result, the correctness tests are expected to fail.

Hope this helps!

fchollet · 2023-07-14T16:53:28Z

Thanks for the PR. It this really necessary? Couldn't you do the same with einsum?

james77777778 · 2023-07-15T14:02:48Z

Thanks for the PR. It this really necessary? Couldn't you do the same with einsum?

Do you mean that affine is not part of the plan, or should I implement the torch's affine operation using einsum?

For 1st:
It is not obvious for me to directly implement affine using einsum.
(actually, I have not seen the implementation of affine in these backends utilizes einsum?)

If the future plan involves implementing KPLs using pure keras-core, the affine op becomes neccesary for many layers.

For 2nd:
I can try to implement torch's map_coordinate by einsum.

fchollet · 2023-07-17T21:58:56Z

What's the closest numpy/scipy op? Does this match https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.affine_transform.html ?

james77777778 · 2023-07-18T01:15:42Z

What's the closest numpy/scipy op? Does this match https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.affine_transform.html ?

The challenge is that the absence of map_coordinates op in TF and Torch. And yes, I want to add affine_transform op to all backends!

In these 3 backends, JAX is the only one that can directly utilize numpy/scipy-like ops because it has map_coordinates op.

For TF, the closest op would be tf.raw_ops.ImageProjectiveTransformV3. This PR mainly follow the signature of tf.raw_ops.ImageProjectiveTransformV3 and the transform definition of [a0, a1, a2, b0, b1, b2, c0, c1]

For Torch, the closest op is a combination of tnn.affine_grid and tnn.grid_sample.

james77777778 · 2023-07-18T03:41:02Z

Usage:

import matplotlib.cbook as cbook
import matplotlib.pyplot as plt
import numpy as np

from keras_core.backend.jax.image import affine as jax_affine
from keras_core.backend.tensorflow.image import affine as tf_affine
from keras_core.backend.torch.image import affine as torch_affine

with cbook.get_sample_data("grace_hopper.jpg") as image_file:
    image = plt.imread(image_file)
transform = np.array([1.2, 0.3, -100, 0.5, 1.2, -150, 0, 0])

jax_affined = jax_affine(image, transform)
tf_affined = tf_affine(image, transform)
torch_affined = torch_affine(image, transform)

fig, ax_dict = plt.subplot_mosaic([["A", "B", "C", "D"]], figsize=(12, 4))
ax_dict["A"].set_title("Original")
ax_dict["A"].imshow(image)
ax_dict["B"].set_title("JAX Affined")
ax_dict["B"].imshow(jax_affined)
ax_dict["C"].set_title("TF Affined")
ax_dict["C"].imshow(tf_affined)
ax_dict["D"].set_title("Torch Affined")
ax_dict["D"].imshow(torch_affined)
fig.tight_layout(h_pad=0.1, w_pad=0.1)
plt.savefig("affine.png")

The visualization:

fchollet · 2023-07-18T04:55:19Z

The challenge is that the absence of map_coordinates op in TF and Torch

@mihirparadkar is working on adding this exact op to Keras -- let's make sure he sees this.

My understanding was that torch's grid_sample was equivalent -- is it not?

Also, if we have map_coordinates, would we still need affine?

james77777778 · 2023-07-18T05:58:25Z

My understanding was that torch's grid_sample was equivalent -- is it not?

I believe grid_sample shares similar functionality but it is not precisely identical in terms of numerical behavior.
(I don't know where the differences come from, but the correctness tests failed.)

Besides:

grid_sample requires normalized coordinates [-1, 1] that is different from map_coordinates.
grid_sample lacks the functionality to fill user-defined value. This PR try to add fill_value by modifying https://github.com/pytorch/vision/blob/main/torchvision/transforms/v2/functional/_geometry.py#L478

Also, if we have map_coordinates, would we still need affine?

As a user, I think it would be more convenient to have affine or affine_transform in keras-core. It is not straightforward to apply affine transform using map_coordinates.
Otherwise there won't exist affine-like apis for these backends:

As a developer, affine is very useful to implement RandomRotation, RandomAffine, etc.

fchollet · 2023-07-18T16:28:43Z

As a user, I think it would be more convenient to have affine or affine_transform in keras-core

Inversely, if we have affine_transform, do we still need map_coordinates?

james77777778 · 2023-07-19T01:13:46Z

Inversely, if we have affine_transform, do we still need map_coordinates?

For common 2D image processing, I believe that affine_transform is sufficient.
However, when it comes to the multidimensional image processing, map_coordinates might be the only option.

A use case for medical image processing:
https://bic-berkeley.github.io/psych-214-fall-2016/map_coordinates.html

fchollet · 2023-07-19T20:58:08Z

Understood. Let's add both affine_transform (same API as scipy, though we don't need to implement all interpolation modes) and map_coordinates. @mihirparadkar is working on map_coordinates.

fchollet

Great work 👍

fchollet · 2023-07-19T20:59:02Z

keras_core/backend/jax/image.py

+)
+
+
+def affine(


Let's call it affine_transform

fchollet · 2023-07-19T20:59:32Z

keras_core/backend/jax/image.py

+    method="bilinear",
+    fill_mode="constant",
+    fill_value=0,
+    data_format="channels_last",


The function signature looks good!

fchollet · 2023-07-19T21:00:19Z

keras_core/backend/jax/image.py

+    if method not in AFFINE_METHODS.keys():
+        raise ValueError(
+            "Invalid value for argument `method`. Expected of one "
+            f"{AFFINE_METHODS.keys()}. Received: method={method}"


This will print "dict_keys()". Let's cast to set() first.

Done for all backends.

fchollet · 2023-07-19T21:00:55Z

keras_core/ops/image.py

+    fill_value=0,
+    data_format="channels_last",
+):
+    # TODO: add docstring


Please add the docstring

Added. I mainly borrowed the docstring from

https://www.tensorflow.org/addons/api_docs/python/tfa/image/transform

https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/image/apply_affine_transform

fchollet · 2023-07-19T21:01:17Z

keras_core/backend/torch/image.py

+    return img
+
+
+def affine(


Are you able to add it for the NumPy backend as well? Fine to depend on scipy.

james77777778 · 2023-07-20T03:06:30Z

Hi @fchollet
I have added affine_transform to NumPy backend and it passed the tests as expected.

Usage:

import matplotlib.cbook as cbook
import matplotlib.pyplot as plt
import numpy as np

from keras_core.backend.jax.image import (
    affine_transform as jax_affine_transform,
)
from keras_core.backend.numpy.image import (
    affine_transform as np_affine_transform,
)
from keras_core.backend.tensorflow.image import (
    affine_transform as tf_affine_transform,
)
from keras_core.backend.torch.image import (
    affine_transform as torch_affine_transform,
)

with cbook.get_sample_data("grace_hopper.jpg") as image_file:
    image = plt.imread(image_file)
transform = np.array([1.2, 0.3, -100, 0.5, 1.2, -150, 0, 0])

jax_affined = jax_affine_transform(image, transform)
np_affined = np_affine_transform(image, transform)
tf_affined = tf_affine_transform(image, transform)
torch_affined = torch_affine_transform(image, transform)

fig, ax_dict = plt.subplot_mosaic([["A", "B", "C", "D", "E"]], figsize=(12, 4))
ax_dict["A"].set_title("Original")
ax_dict["A"].imshow(image)
ax_dict["B"].set_title("JAX Affined")
ax_dict["B"].imshow(jax_affined)
ax_dict["C"].set_title("Numpy Affined")
ax_dict["C"].imshow(np_affined)
ax_dict["D"].set_title("TF Affined")
ax_dict["D"].imshow(tf_affined)
ax_dict["E"].set_title("Torch Affined")
ax_dict["E"].imshow(torch_affined)
fig.tight_layout(h_pad=0.1, w_pad=0.1)
plt.savefig("affine.png")

Visualization:

fchollet

Excellent -- thanks for the polished docstring!

fchollet · 2023-07-20T05:45:55Z

keras_core/ops/image_test.py

+            ("bilinear", "nearest", "channels_last"),
+            ("nearest", "nearest", "channels_last"),
+            # ("bilinear", "wrap", "channels_last"),  no wrap in torch
+            # ("nearest", "wrap", "channels_last"),  no wrap in torch


You can use if backend.backend() == "torch": pytest.skip("wrap fill_mode not support in torch") so that at least it's tested when supported

@fchollet
When using wrap, the output in TF shows inconsistency compared to JAX and NumPy.
What make this situation awkward is that the test sometimes passes but at other times it fails.

As a solution, I have opted to skip the test to prevent CI failures.

The mismatched elements was about 0.6%

Mismatched elements: 45 / 7500 (0.6%) # one run for JAX Mismatched elements: 48 / 7500 (0.64%) # one run for NumPy

fchollet · 2023-07-20T05:46:16Z

keras_core/ops/image.py

+            according to the given mode. Available methods are `"constant"`,
+            `"nearest"` and `"reflect"`. Defaults to `"constant"`.
+        fill_value: Value used for points outside the boundaries of the input if
+            fill_mode=`"constant"`. Defaults to `0`.


Move backticks: fill_mode="constant"

fchollet · 2023-07-20T05:48:40Z

keras_core/ops/image.py

+            gradients are not backpropagated into transformation parameters.
+            Note that `c0` and `c1` are only effective when using TensorFlow
+            backend and will be considered as `0` when using other backends.
+        method: Interpolation method. Available methods are `"nearest"`,


I just noticed that we have a slight API inconsistency: this function as well as resize both use method, but our preprocessing layers use interpolation. Perhaps we should just switch everything (including resize) to interpolation since the preprocessing layers are much older and more widely used.

I have replaced method with interpolation in both affine_transform and resize.
Additionally, there have been some API changes in keras-core.

fchollet

Awesome -- thanks for the great contribution! 👍

* Add affine op * Sync import convention * Use `np.random.random` * Refactor jax implementation * Fix * Address fchollet's comments * Update docstring * Fix test * Replace method with interpolation * Replace method with interpolation * Replace method with interpolation * Update test

james77777778 added 3 commits July 14, 2023 09:03

Add affine op

d95e2dc

Sync import convention

f2f60d8

Use np.random.random

03845ba

james77777778 added 4 commits July 16, 2023 14:40

Merge branch 'keras-team:main' into add-affine

ecc825e

Merge branch 'keras-team:main' into add-affine

a09d8ac

Merge branch 'keras-team:main' into add-affine

70cb67e

Refactor jax implementation

d0ad9b7

Merge branch 'keras-team:main' into add-affine

8d62211

Fix

573af9d

Merge branch 'keras-team:main' into add-affine

029ba49

fchollet reviewed Jul 19, 2023

View reviewed changes

james77777778 added 2 commits July 20, 2023 09:08

Merge branch 'keras-team:main' into add-affine

ea5e916

Address fchollet's comments

fe46aeb

james77777778 requested a review from fchollet July 20, 2023 03:07

Update docstring

7eccadd

james77777778 changed the title ~~Add affine op to all backends~~ Add affine_transform op to all backends Jul 20, 2023

Fix test

031c39e

fchollet reviewed Jul 20, 2023

View reviewed changes

Replace method with interpolation

893089e

james77777778 added 3 commits July 20, 2023 06:57

Replace method with interpolation

ff1db96

Replace method with interpolation

21a6c16

Update test

8794ea0

james77777778 requested a review from fchollet July 20, 2023 07:08

fchollet approved these changes Jul 20, 2023

View reviewed changes

fchollet merged commit 2702048 into keras-team:main Jul 20, 2023

james77777778 deleted the add-affine branch July 21, 2023 02:58

mihirparadkar mentioned this pull request Aug 23, 2023

[WIP] Implements map_coordinates op in keras-core #784

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `affine_transform` op to all backends #477

Add `affine_transform` op to all backends #477

james77777778 commented Jul 14, 2023

fchollet commented Jul 14, 2023

james77777778 commented Jul 15, 2023

fchollet commented Jul 17, 2023

james77777778 commented Jul 18, 2023

james77777778 commented Jul 18, 2023

fchollet commented Jul 18, 2023

james77777778 commented Jul 18, 2023 •

edited

Loading

fchollet commented Jul 18, 2023

james77777778 commented Jul 19, 2023

fchollet commented Jul 19, 2023

fchollet left a comment

fchollet Jul 19, 2023

james77777778 Jul 20, 2023

fchollet Jul 19, 2023

fchollet Jul 19, 2023

james77777778 Jul 20, 2023

fchollet Jul 19, 2023

james77777778 Jul 20, 2023

fchollet Jul 19, 2023

james77777778 Jul 20, 2023

james77777778 commented Jul 20, 2023 •

edited

Loading

fchollet left a comment

fchollet Jul 20, 2023

james77777778 Jul 20, 2023

james77777778 Jul 20, 2023 •

edited

Loading

fchollet Jul 20, 2023

james77777778 Jul 20, 2023

fchollet Jul 20, 2023

james77777778 Jul 20, 2023

fchollet left a comment

		)


		def affine(

Add affine_transform op to all backends #477

Add affine_transform op to all backends #477

Conversation

james77777778 commented Jul 14, 2023

fchollet commented Jul 14, 2023

james77777778 commented Jul 15, 2023

fchollet commented Jul 17, 2023

james77777778 commented Jul 18, 2023

james77777778 commented Jul 18, 2023

fchollet commented Jul 18, 2023

james77777778 commented Jul 18, 2023 • edited Loading

fchollet commented Jul 18, 2023

james77777778 commented Jul 19, 2023

fchollet commented Jul 19, 2023

fchollet left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

james77777778 commented Jul 20, 2023 • edited Loading

fchollet left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

james77777778 Jul 20, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fchollet left a comment

Choose a reason for hiding this comment

Add `affine_transform` op to all backends #477

Add `affine_transform` op to all backends #477

james77777778 commented Jul 18, 2023 •

edited

Loading

james77777778 commented Jul 20, 2023 •

edited

Loading

james77777778 Jul 20, 2023 •

edited

Loading