Make prototype `F` JIT-scriptable #6584

datumbox · 2022-09-14T17:30:43Z

Fixes #6553

co-authored with @vfdev-5

datumbox

Some early notes:

test/test_prototype_transforms_functional.py

torchvision/prototype/transforms/functional/__init__.py

torchvision/prototype/transforms/functional/_meta.py

torchvision/prototype/transforms/functional/_color.py

datumbox · 2022-09-16T14:58:07Z

torchvision/prototype/features/_image.py

@@ -32,6 +32,31 @@ def from_pil_mode(cls, mode: str) -> ColorSpace:
        else:
            return cls.OTHER

+    @staticmethod
+    def from_tensor_shape(shape: List[int]) -> ColorSpace:


I think this makes more sense to be on the Enum than the method. At any case it was moved because JIT didn't like it to be a class method. We don't have to keep it in the enum, we just need to have it as a standalone method.

So we can't even have classmethods on objects that are not covered by JIT? 🤯

Plus, I intentionally put the term "guess" in the name, since the number of channels is not sufficient to pick the right colorspace. For example, CMYK also has 4 channels, but would be classified as RGBA. However, this is not a problem now since we don't support it yet and maybe never will.

I think the issue is that the specific class is definitely not JIT-scriptable. JIT complains about the tensor_contents: Any = None value on the __repr__. Perhaps we can remove this?

I'm happy to make any changes on the name. OR try a different approach with JIT. Let me wrap up the rest of the kernels to see where we are and we can try a couple of options.

If it works as is, I wouldn't mess with this class any further. If it is needed we can remove the : Any though.

@pmeier I'm also really flexible. I don't have context why the tensor_contents is introduced and what's supposed to be. Would you like to send a PR once this is merged to see if you can make it work in the original location?

pmeier

Few comments inline. Otherwise LGTM if CI is green.

test/test_prototype_transforms_functional.py

pmeier · 2022-09-16T15:03:35Z

torchvision/prototype/features/_image.py

@@ -32,6 +32,31 @@ def from_pil_mode(cls, mode: str) -> ColorSpace:
        else:
            return cls.OTHER

+    @staticmethod
+    def from_tensor_shape(shape: List[int]) -> ColorSpace:


So we can't even have classmethods on objects that are not covered by JIT? 🤯

torchvision/prototype/features/_image.py

pmeier · 2022-09-16T15:07:21Z

torchvision/prototype/features/_image.py

@@ -32,6 +32,31 @@ def from_pil_mode(cls, mode: str) -> ColorSpace:
        else:
            return cls.OTHER

+    @staticmethod
+    def from_tensor_shape(shape: List[int]) -> ColorSpace:


Plus, I intentionally put the term "guess" in the name, since the number of channels is not sufficient to pick the right colorspace. For example, CMYK also has 4 channels, but would be classified as RGBA. However, this is not a problem now since we don't support it yet and maybe never will.

torchvision/prototype/transforms/functional/_augment.py

…ecessary defaults.

…ransforms

vfdev-5 · 2022-09-19T15:16:21Z

@datumbox I let you review your PR once you are back :)

datumbox

@vfdev-5 I might not be the best person to "review" this PR but I can confirm that any modification done on your side looks great. It was a good idea to do most of the work needed on fill on separate PRs. I had a look also at your PRs #6595 and #6599 and the direction looks great. Maintaining F JIT-scriptable is a major win, thanks a lot for your work on making it possible.

@pmeier Could you please give us an independent review on this PR?

FYI: follow up work we should do is ensuring the signatures of F and _Feature are aligned. After merging this PR, the signatures of the dispatchers on the base class don't match (one uses Sequences and the other uses List).

datumbox · 2022-09-19T17:02:30Z

torchvision/prototype/transforms/_auto_augment.py

@@ -72,11 +72,10 @@ def _apply_image_transform(

        # Fill = 0 is not equivalent to None, https://github.com/pytorch/vision/issues/6517
        # So, we have to put fill as None if fill == 0
-        fill_: Optional[Union[int, float, Sequence[int], Sequence[float]]]
+        # This is due to BC with stable API which has fill = None by default
+        fill_ = F._geometry._convert_fill_arg(fill)


We might want to consider making _convert_fill_arg() public on the future as this is regularly used in Transforms as a utility method.

vfdev-5 · 2022-09-20T08:49:20Z

FYI: follow up work we should do is ensuring the signatures of F and _Feature are aligned.

@datumbox should they match ? F is now jit scriptable and _Feature is not, so the latter can also accept Sequence[int] or Sequence[float] for user API simplicity:

Image.pad(..., fill=(1, 2, 3))
# instead of only List[float]
Image.pad(..., fill=[1.0, 2.0, 3.0])

Do we want that ? If yes, I can update this PR and put all _Feature fill typehint as Optional[Union[int, float, List[float]]] = None

datumbox · 2022-09-20T08:58:10Z

@vfdev-5 One of the original proposals was that the _Feature class and F will have methods with aligned signatures. The first would receive the maximal possible parameters (even those not used) and do a simple redirection to the low-level kernel. The only reason I proposed to change them is to ensure this is uphold. We can certainly do this on a follow up PR where we confirm that all APIs are aligned overall. What do you think?

vfdev-5 · 2022-09-20T09:09:02Z

@datumbox OK to me. @pmeier please review and we merge this PR and I'll update _Feature signatures in a follow-up PR

pmeier

Mostly questions from my side. LGTM if CI is green. Thanks @datumbox and @vfdev-5!

pmeier · 2022-09-20T11:01:18Z

test/prototype_transforms_dispatcher_infos.py

+    # Enum<__torch__.torchvision.transforms.functional.InterpolationMode> interpolation=Enum<InterpolationMode.BILINEAR>,
+    # Union(float[], float, int, NoneType) fill=None) -> Tensor
+    #
+    # This is probably due to the fact that F.perspective does not have the same signature as F.perspective_image_tensor


Yes, if the signatures actually diverge, then this is the issue? I'm ok disabling this here for now and look at it later.

Still, is this something we want? Shouldn't the public kernels be in sync with the dispatcher?

I believe @vfdev-5 made this change. As discussed previously there are other places where we need to align the signatures and this can happen on a follow up PR to avoid making this too long.

Personally I think it's worth aligning the signatures unless there is a good reason not to (perhaps the interpolation default value is one exception?). I'm open to discussing this and I think we should agree on the policy soon.

pmeier · 2022-09-20T11:01:55Z

test/test_prototype_transforms_dispatchers.py

@@ -9,7 +9,6 @@


 class TestCommon:
-    @pytest.mark.xfail(reason="dispatchers are currently not scriptable")


pmeier · 2022-09-20T11:02:59Z

test/test_prototype_transforms_functional.py

+def test_scriptable_midlevel(midlevel):
+    jit.script(midlevel)  # TODO: pass data through it


@vfdev-5 Isn't that superseded by the tests I've added?

Maybe? I noticed that at one point early in the development, one of the kernels had a return type Any and this test didn't complaint. This is why usually to be safe, we take the strategy of:

JIT-scripting

Passing data through the kernel and compare it with the non-JIT version

Serialize/deserialize the method and confirm it still returns the right value.

See _check_jit_scriptable() from test_models.py for more info.

@datumbox in test/test_prototype_transforms_dispatchers.py the test should script and execute midlevel op, so I think we can remove this test_scriptable_midlevel

I'll cleanup.

torchvision/prototype/features/_feature.py

pmeier · 2022-09-20T11:05:24Z

torchvision/prototype/features/_image.py

@@ -32,6 +32,31 @@ def from_pil_mode(cls, mode: str) -> ColorSpace:
        else:
            return cls.OTHER

+    @staticmethod
+    def from_tensor_shape(shape: List[int]) -> ColorSpace:


If it works as is, I wouldn't mess with this class any further. If it is needed we can remove the : Any though.

pmeier · 2022-09-20T11:08:41Z

torchvision/prototype/transforms/_geometry.py

@@ -252,7 +252,14 @@ def __init__(

    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
        fill = self.fill[type(inpt)]
-        return F.pad(inpt, padding=self.padding, fill=fill, padding_mode=self.padding_mode)
+
+        # This cast does Sequence[int] -> List[int] and is required to make mypy happy


Well, mypy is not at fault here. The kernel could support Sequence[int], but we annotated List[int]. Since that is more specific, mypy is right to complain. We can always silence mypy with a # type: ignore[...] comment or use cast(...) since both have no runtime implications. Actually casting to a list while it is not needed hurts performance although probably not by a significant amount.

At this point where we try to improve the speed by any margin, I would prefer using ignore statements rather than casting. I don't know how negligible the introduction of cast() is but if there is no measurable slow-down I don't mind switching.

Summary: * Improve existing low kernel test. * Add new midlevel jit-scriptability test (failing). * Remove duplicate aliases from kernel tests. * Fixing colour kernels. * Fixing deprecated kernels. * fix mypy * Silence mypy instead of fixing to avoid performance penalty * Fixing augment kernels. * Fixing augment meta. * Remove is_tracing calls. * Add fake ImageType and DType * Fixing type conversion kernels. * Fixing misc kernels. * partial fix geometry * Remove mutable default from `_pad_with_vector_fill()` + all other unnecessary defaults. * Fix geometry ops * Fixing tests * Removed xfail for jit tests on midlevel ops Reviewed By: NicolasHug Differential Revision: D39765297 fbshipit-source-id: 50ec9dc9d9e2f9c8dab6ab01337e01643dc0ab64 Co-authored-by: vfdev-5 <[email protected]>

datumbox added 2 commits September 14, 2022 18:28

Improve existing low kernel test.

136da06

Add new midlevel jit-scriptability test (failing).

e5dfa68

facebook-github-bot added the cla signed label Sep 14, 2022

datumbox marked this pull request as draft September 14, 2022 17:30

datumbox mentioned this pull request Sep 14, 2022

Rename convert_bounding_box_format => convert_format_bounding_box #6582

Merged

Merge branch 'main' into jit/prototype_transforms

6624ab5

datumbox mentioned this pull request Sep 16, 2022

Clean up prototype transforms before migration #5626

Open

Remove duplicate aliases from kernel tests.

9369c6e

datumbox commented Sep 16, 2022

View reviewed changes

datumbox added 2 commits September 16, 2022 15:26

Fixing colour kernels.

087e916

Fixing deprecated kernels.

f844a5c

datumbox commented Sep 16, 2022

View reviewed changes

pmeier approved these changes Sep 16, 2022

View reviewed changes

datumbox requested a review from pmeier September 16, 2022 15:24

datumbox added 3 commits September 16, 2022 16:49

fix mypy

0b961ef

Silence mypy instead of fixing to avoid performance penalty

f75a0df

Fixing augment kernels.

c286621

vfdev-5 reviewed Sep 16, 2022

View reviewed changes

torchvision/prototype/transforms/functional/_augment.py Outdated Show resolved Hide resolved

datumbox added 5 commits September 16, 2022 17:22

Fixing augment meta.

1c2822c

Remove is_tracing calls.

4a3df83

Add fake ImageType and DType

663d8c7

Fixing type conversion kernels.

0d38972

Fixing misc kernels.

60000fa

datumbox force-pushed the jit/prototype_transforms branch from ef7130c to 60000fa Compare September 16, 2022 17:26

datumbox and others added 2 commits September 16, 2022 19:06

partial fix geometry

58e5707

Merge branch 'main' into jit/prototype_transforms

1e301b1

datumbox force-pushed the jit/prototype_transforms branch from aafb356 to 1e301b1 Compare September 17, 2022 08:39

datumbox and others added 2 commits September 17, 2022 09:48

Merge branch 'main' into jit/prototype_transforms

2f37578

Remove mutable default from _pad_with_vector_fill() + all other unn…

1395149

…ecessary defaults.

datumbox mentioned this pull request Sep 17, 2022

Updated fill arg typehint for affine, perspective and elastic ops #6595

Merged

datumbox mentioned this pull request Sep 17, 2022

[proto] Added dict support for fill arg for remaining transforms #6599

Merged

vfdev-5 added 6 commits September 19, 2022 11:42

Merge branch 'main' of github.com:pytorch/vision into jit/prototype_t…

96f318d

…ransforms

Fix geometry ops

87f3567

Fixing tests

df53707

Merge branch 'main' of github.com:pytorch/vision into jit/prototype_t…

e0185fa

…ransforms

Merge branch 'main' of github.com:pytorch/vision into jit/prototype_t…

3d0e1c9

…ransforms

Removed xfail for jit tests on midlevel ops

d7e516c

vfdev-5 marked this pull request as ready for review September 19, 2022 15:15

vfdev-5 changed the title ~~[WIP] Make prototype F JIT-scriptable~~ Make prototype F JIT-scriptable Sep 19, 2022

datumbox commented Sep 19, 2022

View reviewed changes

pmeier approved these changes Sep 20, 2022

View reviewed changes

datumbox added enhancement module: transforms prototype labels Sep 20, 2022

Merge branch 'main' into jit/prototype_transforms

7c309eb

datumbox merged commit 841b9a1 into pytorch:main Sep 20, 2022

datumbox deleted the jit/prototype_transforms branch September 20, 2022 11:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make prototype `F` JIT-scriptable #6584

Make prototype `F` JIT-scriptable #6584

datumbox commented Sep 14, 2022 •

edited

Loading

datumbox left a comment

datumbox Sep 16, 2022

pmeier Sep 16, 2022

pmeier Sep 16, 2022

datumbox Sep 16, 2022

pmeier Sep 20, 2022

datumbox Sep 20, 2022

pmeier left a comment

pmeier Sep 16, 2022

pmeier Sep 16, 2022

vfdev-5 commented Sep 19, 2022

datumbox left a comment •

edited

Loading

datumbox Sep 19, 2022

vfdev-5 commented Sep 20, 2022

datumbox commented Sep 20, 2022 •

edited

Loading

vfdev-5 commented Sep 20, 2022 •

edited

Loading

pmeier left a comment

pmeier Sep 20, 2022

datumbox Sep 20, 2022

pmeier Sep 20, 2022

pmeier Sep 20, 2022

datumbox Sep 20, 2022

vfdev-5 Sep 20, 2022

pmeier Sep 20, 2022

pmeier Sep 20, 2022

pmeier Sep 20, 2022

datumbox Sep 20, 2022

		@@ -9,7 +9,6 @@


		class TestCommon:
		@pytest.mark.xfail(reason="dispatchers are currently not scriptable")

		def test_scriptable_midlevel(midlevel):
		jit.script(midlevel) # TODO: pass data through it

Make prototype F JIT-scriptable #6584

Make prototype F JIT-scriptable #6584

Conversation

datumbox commented Sep 14, 2022 • edited Loading

datumbox 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

pmeier left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vfdev-5 commented Sep 19, 2022

datumbox left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vfdev-5 commented Sep 20, 2022

datumbox commented Sep 20, 2022 • edited Loading

vfdev-5 commented Sep 20, 2022 • edited Loading

pmeier 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

Choose a reason for hiding this comment

Make prototype `F` JIT-scriptable #6584

Make prototype `F` JIT-scriptable #6584

datumbox commented Sep 14, 2022 •

edited

Loading

datumbox left a comment •

edited

Loading

datumbox commented Sep 20, 2022 •

edited

Loading

vfdev-5 commented Sep 20, 2022 •

edited

Loading