Add PixelUnshuffle #49334

jbschlosser · 2020-12-14T16:28:58Z

Summary: Adds an implementation of torch.nn.PixelUnshuffle as the inverse operation of torch.nn.PixelShuffle. This addresses #2456

Test Plan:

# Unit tests.
python test/test_nn.py TestNN.test_pixel_shuffle_unshuffle

# Module test.
python test/test_nn.py TestNN.test_PixelUnshuffle

# C++ API tests.
build/bin/test_api

# C++ / python parity tests.
python test/test_cpp_api_parity.py

# JIT test.
python test/test_jit.py TestJitGeneratedFunctional.test_nn_pixel_unshuffle

# Override tests.
python test/test_overrides.py

# Type hint tests.
python test/test_type_hints.py

Differential Revision: D25401439

Screenshots of rendered docs:

facebook-github-bot · 2020-12-15T02:09:42Z

This pull request was exported from Phabricator. Differential Revision: D25401439

gchanan · 2020-12-15T20:48:00Z

aten/src/ATen/native/PixelShuffle.cpp

+  int64_t h = self.size(-2);
+  int64_t w = self.size(-1);
+  const auto NUM_NON_BATCH_DIMS = 3;
+  const auto last_batch_dim = self.sizes().end() - NUM_NON_BATCH_DIMS;


nit: I would definitely expect something called last_batch_dim to be an integer type, that it's only valid on self_sizes should probably be part of the name.

also, doesn't this point to the first non-batch dim? I.e. it's only the last_batch_dim because you are calling it in a function that "ends" one past. So maybe a better name is "self_sizes_batch_end" or something.

Yes, you're right- it points to just after the last batch dim. To be more precise, I'll change this to self_sizes_batch_end as suggested.

gchanan · 2020-12-15T20:56:10Z

aten/src/ATen/native/PixelShuffle.cpp

+  // (ow, downscale_factor) dims. This allows unshuffling to be done next by permuting dims.
+  std::vector<int64_t> expanded_shape(self.sizes().begin(), last_batch_dim);
+  expanded_shape.insert(expanded_shape.end(), {c, oh, downscale_factor, ow, downscale_factor});
+  const auto input_expanded = self.reshape(expanded_shape);


"expand" suggests something particular, which is that it called "expand" (i.e. didn't allocate new memory). Is that what you expect here?

Good point- expand as in tensor.expand() is not what is intended here. Ideally, I want to signal that the input has been reshaped, resulting in an increased number of input dimensions with no change to the number of elements. Would input_reshaped work here or is there a more specific way to name this? Any suggestions on the name for the shape itself?

gchanan · 2020-12-15T20:59:22Z

torch/nn/functional.py

+pixel_unshuffle = _add_docstr(torch.pixel_unshuffle, r"""
+pixel_unshuffle(input, downscale_factor) -> Tensor
+
+Rearranges elements in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a


do we use * in the docs consistently how you've used it here (i.e. correctly?). I'm particularly asking about the case where * translates to no dimensions.

Hmm good question! From the nn.Linear docs:

"Input: (N, *, H_{in}), where * means any number of additional dimensions and H_{in} = in_features"

This description indicates to me that N and H_{in} are required dimensions, and * can be any number of additional dimensions in between those two. However, surprisingly to me, nn.Linear does accept an input shape of (H_{in},). It seems that the * functions like "0 or more of the previous" (i.e. the Kleene star). I think it is the use of the word "additional" that makes this surprising to me, but maybe this is only a problem for me.

Similarly in nn.L1Loss:

"Input: (N, *) where * means, any number of additional dimensions"

nn.L1Loss does indeed accept tensors of 0 dimensions, so again the * functions like a Kleene star.

TL;DR: To match what's been done in the two referenced modules, the line here should change to: (N, *, C, H \times r, W \times r). Similar updates should be done in PixelShuffle for consistency.

We do typically use * for 0 or more dimensions. See, for example, the torch.svd documentation: https://pytorch.org/docs/master/generated/torch.svd.html?highlight=svd#torch.svd. We're probably inconsistent about this in our documentation.

I could see N* or * as appropriate. I wouldn't want to use (N, *, C...) unless zero or more dimensions could be inserted between a required dimension N and a required dimension C.

I could see N* or * as appropriate. I wouldn't want to use (N, *, C...) unless zero or more dimensions could be inserted between a required dimension N and a required dimension C.

I agree with this as it matches my intuition, and I think we should make a note to eventually make this consistent across torch.nn modules (including nn.Linear and nn.L1Loss).

As far as this PR is concerned, I will add a note to the docs indicating what * represents here, with similar phrasing to that of torch.svd.

mruberry · 2020-12-16T17:51:01Z

torch/nn/functional.py

+pixel_unshuffle(input, downscale_factor) -> Tensor
+
+Rearranges elements in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a
+tensor of shape :math:`(*, C \times r^2, H, W)`.


Like the pixel shuffle documentation, this references r before it's defined. The input argument is "downscale_factor". Also, while this first statement is true, I think the highest order bit to communicate when describing this function is that it reverses the pixel shuffle operation, which can can be linked to in the first sentence. The identify can even be elaborated on in a mathematical line showing that unshuffle(shuffle(...)) is an identity. This first sentence is OK, and the reshaping performed is important, but it loses the context of what the reshape does, and there's no mention of how the values in the tensor are permuted. That is, an extremely cynical reading of this first sentence is simply that it randomly rearranges all the elements in the tensor and puts them into the new shape.

mruberry · 2020-12-16T17:51:50Z

torch/nn/functional.py

+Rearranges elements in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a
+tensor of shape :math:`(*, C \times r^2, H, W)`.
+
+See :class:`~torch.nn.PixelUnshuffle` for details.


It always seems like an interesting quirk to me that we refer to the module over the function as our source of truth, but I think this is how most of our nn documents work.

mruberry · 2020-12-16T17:52:30Z

torch/nn/modules/pixelshuffle.py

+
+
+class PixelUnshuffle(Module):
+    r"""Rearranges elements in a tensor of shape :math:`(*, C, H \times r, W \times r)`


Same comments as for above here.

mruberry · 2020-12-16T17:53:20Z

torch/nn/modules/pixelshuffle.py

+    to a tensor of shape :math:`(*, C \times r^2, H, W)`. This is the inverse operation
+    of :class:`~torch.nn.PixelShuffle`.
+
+    Note that this function can take inputs with any number of batch dimensions:


I don't think we need this note and should reinforce that * is zero or more batch dimensions, since that pattern appears consistently in the docs.

torch/nn/modules/pixelshuffle.py

mruberry · 2020-12-16T17:56:55Z

aten/src/ATen/native/PixelShuffle.cpp

@@ -51,4 +51,49 @@ Tensor pixel_shuffle(const Tensor& self, int64_t upscale_factor) {
  return input_permuted.reshape(final_shape);
 }

+
+Tensor pixel_unshuffle(const Tensor& self, int64_t downscale_factor) {
+  TORCH_CHECK(self.dim() >= 3,


Should there also be a check that downscale factor is > 0?

I like this, especially because the below checks of h % downscale_factor == 0 and w % downscale_factor == 0 may pass even with a negative downscale_factor due to the specific way mod is implemented in C++.

For consistency, what do you think about me adding an upscale_factor check in pixel_shuffle in this PR?

glaringlee · 2020-12-18T05:32:23Z

torch/nn/functional.py

+pixel_unshuffle = _add_docstr(torch.pixel_unshuffle, r"""
+pixel_unshuffle(input, downscale_factor) -> Tensor
+
+Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape :math:`(*, C \times r^2, H, W)`, where r is a downscale factor.


This might cause lint problem.
max number of character per line is 120.

Thanks! Just setup my local dev environment and the python linter (flake8) is not functioning correctly yet :(

glaringlee · 2020-12-18T05:33:05Z

torch/nn/modules/pixelshuffle.py

+
+
+class PixelUnshuffle(Module):
+    r"""Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape :math:`(*, C \times r^2, H, W)`, where r is a downscale factor.


same here.
This might cause lint problem.
max number of character per line is 120.

mruberry · 2020-12-20T09:08:24Z

I reviewed the test failures, most are:

Dec 18 08:04:54 AssertionError: The following functions are not tested for __torch_function__ support, please ensure there is an entry in the dict returned by torch._overrides.get_testing_overrides for this function or if a __torch_function__ override does not make sense, add an entry to the tuple returned by torch._overrides.get_ignored_functions.

Is caused because the function is missing an override entry. See @soulitzer's recent PR adding torch.sinc for an idea of what files need to be updated: #48740. The relevant line is this one:

pytorch/torch/overrides.py

Line 774 in c0deb23

torch.sinc: lambda input, out=None: -1,

@heitorschueroff has also been interested in helping improve our documentation for adding a new operator. There's a lot to it.

A few builds are also failing due to a couple mypy issues, and rebasing is likely to fix the rest.

Summary: Pull Request resolved: pytorch#49334 Adds an implementation of `torch.nn.PixelUnshuffle` as the inverse operation of `torch.nn.PixelShuffle`. This addresses pytorch#2456 Test Plan: `buck test caffe2/test:nn -- test_pixel_unshuffle` Differential Revision: D25401439 fbshipit-source-id: edf92f0ac884410c287d7afc91506d50b52597cd

codecov · 2020-12-22T06:50:09Z

Codecov Report

Merging #49334 (26c3b7a) into master (7ed140a) will increase coverage by 5.32%.
The diff coverage is 97.22%.

@@            Coverage Diff             @@
##           master   #49334      +/-   ##
==========================================
+ Coverage   75.24%   80.56%   +5.32%     
==========================================
  Files        1883     1887       +4     
  Lines      204470   204663     +193     
==========================================
+ Hits       153851   164885   +11034     
+ Misses      50619    39778   -10841

facebook-github-bot

@jbschlosser has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

mruberry · 2020-12-22T17:12:13Z

aten/src/ATen/native/PixelShuffle.cpp

+  int64_t c = self.size(-3);
+  int64_t h = self.size(-2);
+  int64_t w = self.size(-1);
+  const auto NUM_NON_BATCH_DIMS = 3;


nit: constexpr

mruberry · 2020-12-22T17:23:33Z

torch/nn/functional.py

@@ -2799,7 +2799,7 @@ def multi_margin_loss(input, target, p=1, margin=1., weight=None, size_average=N
 pixel_shuffle(input, upscale_factor) -> Tensor

 Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)` to a
-tensor of shape :math:`(*, C, H \times r, W \times r)`.
+tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor.


"r is the :attr:`upscale_factor`."

mruberry · 2020-12-22T17:23:55Z

torch/nn/functional.py

+
+Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a
+tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape
+:math:`(*, C \times r^2, H, W)`, where r is a downscale factor.


Analogous change here as above.

mruberry · 2020-12-22T17:26:47Z

torch/nn/modules/pixelshuffle.py

@@ -6,26 +6,30 @@

 class PixelShuffle(Module):
    r"""Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)`
-    to a tensor of shape :math:`(*, C, H \times r, W \times r)`.
+    to a tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor.


Using "an upscale factor" seems OK here since the name of the input is not available

mruberry

LGTM! Just some minor inline comments.

glaringlee

@jbschlosser
This generally looks good to me now, I have one nit comment, not mandatory.

glaringlee · 2020-12-22T18:46:00Z

torch/csrc/api/include/torch/nn/options/pixelshuffle.h

+/// PixelUnshuffle model(PixelUnshuffleOptions(5));
+/// ```
+struct TORCH_API PixelUnshuffleOptions {
+  PixelUnshuffleOptions(int64_t downscale_factor)


nit: for single arg unique constructor, we normally put explicit before, this can avoid warnings in phabricator as well.

Thanks for catching this! I might be wrong, but it looks like the implicit behavior is desired for C++ module options. None of the option structs defined in torch/csrc/api/include/torch/nn/options use explicit, and in fact several have the /* implicit */ comment. Further, the C++ module tests are often written to use the implicit conversion. So to follow our standards and match these, I think I should add a similar /* implicit */ comment here as well. Thoughts on this?

@jbschlosser
The C++ lib is contributed by many external contributors. Multiple coding styles are mixed together at the moment.
I plan to standardized it a bit in the future.
Adding /*implicit*/ is fine, I will update them in one shot later.

facebook-github-bot

@jbschlosser has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2020-12-23T05:12:31Z

@jbschlosser merged this pull request in 68d438c.

facebook-github-bot added fb-exported cla signed labels Dec 14, 2020

jbschlosser force-pushed the export-D25401439 branch from 508bddc to d5aa90c Compare December 14, 2020 16:34

jbschlosser force-pushed the export-D25401439 branch from d5aa90c to cbb0615 Compare December 15, 2020 02:09

jbschlosser changed the title ~~[pytorch] Add PixelUnshuffle~~ Add PixelUnshuffle Dec 15, 2020

gchanan reviewed Dec 15, 2020

View reviewed changes

mruberry reviewed Dec 16, 2020

View reviewed changes

torch/nn/modules/pixelshuffle.py Show resolved Hide resolved

mruberry reviewed Dec 16, 2020

View reviewed changes

jbschlosser requested a review from glaringlee as a code owner December 17, 2020 22:42

glaringlee reviewed Dec 18, 2020

View reviewed changes

jbschlosser added 7 commits December 21, 2020 10:29

Addressing PR comments (doc changes + extra checks + renaming)

bd143a4

PixelUnshuffle: C++ parity & tests + lint + better py tests

1bdd60c

Linter strikes again

e02d893

Manually fix long lines (need to get flake8 running properly)

3ea84a5

Adding override func and interned string for pixel unshuffle

07c33ed

Adding JIT test & gen_pyi entry for pixel unshuffle

cbe6c1c

jbschlosser force-pushed the export-D25401439 branch from ad17f9f to cbe6c1c Compare December 21, 2020 17:27

Fix doc warning: horizontal line too short

f6cda98

jbschlosser requested review from mruberry, gchanan and glaringlee December 22, 2020 14:44

facebook-github-bot reviewed Dec 22, 2020

View reviewed changes

mruberry reviewed Dec 22, 2020

View reviewed changes

mruberry self-requested a review December 22, 2020 17:28

mruberry approved these changes Dec 22, 2020

View reviewed changes

Addressing PR comments

7b0d728

glaringlee reviewed Dec 22, 2020

View reviewed changes

Adding implicit comment to PixelUnshuffleOptions

26c3b7a

facebook-github-bot reviewed Dec 22, 2020

View reviewed changes

facebook-github-bot closed this in 68d438c Dec 23, 2020

facebook-github-bot added the Merged label Dec 23, 2020

jbschlosser mentioned this pull request Mar 19, 2021

PixelShuffle for downscaling. (Feature proposal) #1684

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add PixelUnshuffle #49334

Add PixelUnshuffle #49334

jbschlosser commented Dec 14, 2020 •

edited

facebook-github-bot commented Dec 15, 2020

gchanan Dec 15, 2020

gchanan Dec 15, 2020

jbschlosser Dec 15, 2020

gchanan Dec 15, 2020

jbschlosser Dec 15, 2020

gchanan Dec 15, 2020

jbschlosser Dec 15, 2020

mruberry Dec 16, 2020

jbschlosser Dec 16, 2020

jbschlosser Dec 16, 2020

mruberry Dec 16, 2020

mruberry Dec 16, 2020

mruberry Dec 16, 2020

mruberry Dec 16, 2020

mruberry Dec 16, 2020

jbschlosser Dec 16, 2020

glaringlee Dec 18, 2020

jbschlosser Dec 18, 2020

glaringlee Dec 18, 2020

mruberry commented Dec 20, 2020 •

edited

codecov bot commented Dec 22, 2020 •

edited

facebook-github-bot left a comment

mruberry Dec 22, 2020

mruberry Dec 22, 2020

mruberry Dec 22, 2020

mruberry Dec 22, 2020

mruberry left a comment

glaringlee left a comment

glaringlee Dec 22, 2020

jbschlosser Dec 22, 2020

glaringlee Dec 22, 2020 •

edited

facebook-github-bot left a comment

facebook-github-bot commented Dec 23, 2020



		class PixelUnshuffle(Module):
		r"""Rearranges elements in a tensor of shape :math:`(*, C, H \times r, W \times r)`



		class PixelUnshuffle(Module):
		r"""Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a tensor of shape :math:`(, C, H \times r, W \times r)` to a tensor of shape :math:`(, C \times r^2, H, W)`, where r is a downscale factor.

Add PixelUnshuffle #49334

Add PixelUnshuffle #49334

Conversation

jbschlosser commented Dec 14, 2020 • edited

facebook-github-bot commented Dec 15, 2020

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

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

mruberry commented Dec 20, 2020 • edited

codecov bot commented Dec 22, 2020 • edited

Codecov Report

facebook-github-bot 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

mruberry left a comment

Choose a reason for hiding this comment

glaringlee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glaringlee Dec 22, 2020 • edited

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot commented Dec 23, 2020

jbschlosser commented Dec 14, 2020 •

edited

mruberry commented Dec 20, 2020 •

edited

codecov bot commented Dec 22, 2020 •

edited

glaringlee Dec 22, 2020 •

edited