Repository: pytorch/vision
Pull request: Update docstring for RandAugment
Fixes #4456 

Problem statement:
AutoAugment's document is wrong (or unexpected?)
### 🐛 Describe the bug

In the document, it is said that:
> [...] If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have […, 1 or 3, H, W] shape, where … means an arbitrary number of leading dimensions.

So, if I understand correctly, the shape `(1,H,W)` still works?

However, when I pass an input tensor whose shape is `(1,H,W)`, it sometimes does the job and sometimes doesn't.

To reproduce:

First, since it is hard to reproduce a randomized transformation that I have just suddenly met, I have to loop through several random seeds using the code below:

```python
for seed in range(0, 1000):
    print(seed)
    torch.manual_seed(seed)
    image = torch.randint(0,255,(1,32,64), dtype=torch.uint8)
    tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
    image = tf(image)
    image
```

And the error raises when `seed=5`, I then hard-code the seed as follows:
```python
torch.manual_seed(5)
image = torch.randint(0,255,(1,32,64), dtype=torch.uint8)
tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
image = tf(image)
image
```

The traceback error is:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_32333/3324237737.py in <module>
      2 image = torch.randint(0,255,(1,1,32,64), dtype=torch.uint8)
      3 tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
----> 4 image = tf(image)
      5 image

~/miniconda3/envs/cv/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/autoaugment.py in forward(self, img)
    214                     img = F.adjust_brightness(img, 1.0 + magnitude)
    215                 elif op_name == "Color":
--> 216                     img = F.adjust_saturation(img, 1.0 + magnitude)
    217                 elif op_name == "Contrast":
    218                     img = F.adjust_contrast(img, 1.0 + magnitude)

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional.py in adjust_saturation(img, saturation_factor)
    810         return F_pil.adjust_saturation(img, saturation_factor)
    811 
--> 812     return F_t.adjust_saturation(img, saturation_factor)
    813 
    814 

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py in adjust_saturation(img, saturation_factor)
    213     _assert_image_tensor(img)
    214 
--> 215     _assert_channels(img, [3])
    216 
    217     return _blend(img, rgb_to_grayscale(img), saturation_factor)

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py in _assert_channels(img, permitted)
     53     c = _get_image_num_channels(img)
     54     if c not in permitted:
---> 55         raise TypeError("Input image tensor permitted channel values are {}, but found {}".format(permitted, c))
     56 
     57 

TypeError: Input image tensor permitted channel values are [3], but found 1
```

I think some operators do not work (or it hasn't been designed to work) with grayscale images?
I still randomly meet this error with IMAGENET policy also.

### Versions

```
Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.27

Python version: 3.8.11 (default, Aug  3 2021, 15:09:35)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.0.0-23-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2060
Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h497a2fe_0    conda-forge
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h1abd341_0    conda-forge
[conda] numpy                     1.20.3           py38hf144106_0  
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch                   1.9.0           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchvision               0.10.0               py38_cu102    pytorch
```

cc @vfdev-5 @datumbox

Hints from issue discussion:
@VinhLoiIT You are absolutely right, the documentation is wrong. Only RGB images are supported.

The reason this works sometimes is due to the randomness of AA, which leads to choosing transforms that can be applied on non RGB images. But then when a transformation like `adjust_saturation` is sampled that requires RGB, it fails.

The doc string is wrong on all 3 autoaugment methods:
https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L68

https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L252

https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L336

We should be saying that only `[..., 3, H, W]` shape is supported. Would you be interested in sending a PR that fixes the problem?

Commit message: update docstring for RandAugment

Patch to review:
diff --git a/torchvision/transforms/autoaugment.py b/torchvision/transforms/autoaugment.py
index bffc4a24f67..257fe538711 100644
--- a/torchvision/transforms/autoaugment.py
+++ b/torchvision/transforms/autoaugment.py
@@ -65,7 +65,7 @@ class AutoAugment(torch.nn.Module):
     r"""AutoAugment data augmentation method based on
     `"AutoAugment: Learning Augmentation Strategies from Data" <https://arxiv.org/pdf/1805.09501.pdf>`_.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
     If img is PIL Image, it is expected to be in mode "L" or "RGB".
 
     Args:
@@ -249,7 +249,7 @@ class RandAugment(torch.nn.Module):
     `"RandAugment: Practical automated data augmentation with a reduced search space"
     <https://arxiv.org/abs/1909.13719>`_.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
     If img is PIL Image, it is expected to be in mode "L" or "RGB".
 
     Args:
@@ -333,7 +333,7 @@ class TrivialAugmentWide(torch.nn.Module):
     r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
     `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
     If img is PIL Image, it is expected to be in mode "L" or "RGB".
 
     Args:


Code review comments:
diff_hunk: @@ -249,7 +249,7 @@ class RandAugment(torch.nn.Module):
     `"RandAugment: Practical automated data augmentation with a reduced search space"
     <https://arxiv.org/abs/1909.13719>`_.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
     If img is PIL Image, it is expected to be in mode "L" or "RGB".
path: torchvision/transforms/autoaugment.py
text: @user1:
```suggestion
    If img is PIL Image, it is expected to be in mode "RGB".
```

@author:
Oh, I forgot about changing the pillow part also. Thanks for pointing it.

diff_hunk: @@ -333,7 +333,7 @@ class TrivialAugmentWide(torch.nn.Module):
     r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
     `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
     If img is PIL Image, it is expected to be in mode "L" or "RGB".
path: torchvision/transforms/autoaugment.py
text: @user1:
```suggestion
    If img is PIL Image, it is expected to be in mode "RGB".
```

@author:
Hi @user1 ,

The below snippet still works when I pass a pillow image with the mode is `'L'`. This definition of "work" here is that it does not raise the above issue. I've also checked some of the results and see that they are all good. If then, does it come from the tensor transformation side?

```python
from PIL import Image
from pathlib import Path
import torch
import torchvision
from torchvision import transforms

image_path = 'n02100877_4699.jpg'
raw_image = Image.open(image_path).convert('L')  # try with 'L', 'RGB
print(image.size, image.mode) 
tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
for seed in range(0, 2000):
    torch.manual_seed(seed)
    image = tf(raw_image)
    image.save(out_dir / f"{seed}.png", 'png')
```

@author:
Here is the already setup colab notebook that you could manually check yourself: https://colab.research.google.com/drive/1WtlOAASAGJA927oeFECN2YSFAhlMJ6j2?usp=sharing

@user1:
Hmm, thanks for investigating. It seems that PIL supports the colour operators even when the image is grayscale. 

I think we have a few options:
1. Update the documentation for Tensors only since PIL supports `L`
2. Update the documentation for both to avoid confusing people
3. Update the implementations of colour transforms of TorchVision that don't support Grayscale images to align them with PIL.

I think option 3 is worth exploring but it will take longer and needs to happen carefully. So I would be in favour of doing 1 or 2 in this PR and follow up on 3 on a separate one. I don't have strong opinion over how we go, so let me know what you prefer.

@author:
Yeah, I think that this PR is just for updating the documentation only to keep it relates to the original issue. Furthermore, since we have already declared some policies such as ImageNet, CIFAR10 that their input is not a grayscale image. Thus, changing to support grayscale images would have some conflicts with the policy definitions as well that we might have further discussion in another thread.

In conclusion, in this PR, I would change the document so that this transformation only supports RGB for now.