PyTorch JIT trace on Swin Transformer pretrained checkpoint fails
Author: hollance
Pull request: no
Labels: name: Good Difficult Issue

name: bug

### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.9.12
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.11.0 (False)
- Tensorflow version (GPU?): 2.8.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.4.1 (cpu)
- Jax version: 0.3.7
- JaxLib version: 0.3.7
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no
```


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

The goal is to run `torch.jit.trace` on the Swin Transformer model from a pretrained checkpoint so it can be exported to another format (e.g. Core ML, ONNX, etc).

Steps to reproduce the issue are:

```python
from transformers import SwinForImageClassification
model_checkpoint = "microsoft/swin-small-patch4-window7-224"
model = SwinForImageClassification.from_pretrained(model_checkpoint, torchscript=True).eval()

import torch
example_input = torch.rand([1, model.config.num_channels, model.config.image_size, model.config.image_size])

traced_model = torch.jit.trace(model, example_input, strict=False)
```

The trace gives a lot of warnings, which are not important, and then fails with the error:

```python
TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!
```

### Expected behavior

There should be no error and the trace completes successfully.

The unit tests for Swin Transformer also perform a JIT trace and they succeed. The model architecture used in the test case is simpler than from the pretrained checkpoint.

In particular, the issue seems to be the `window_size` parameter. This can be demonstrated as follows:

```python
from transformers import SwinConfig
config = SwinConfig.from_pretrained(model_checkpoint, window_size=6)
config.torchscript = True

model = SwinForImageClassification.from_pretrained(model_checkpoint, config=config, ignore_mismatched_sizes=True).eval()

traced_model = torch.jit.trace(model, example_input, strict=False)
```

By forcing the window size to be a different number, the trace now completes without errors. (But of course, this window size is not appropriate for the pretrained checkpoint.)



Comments:
Using `check_trace=False` with `torch.jit.trace` avoids this problem. The traced model seems to work OK and gives the same results as the original (tested with several different input images).

It's probably still a good idea to make the trace work without the error (if possible) but as using `check_trace=False` solves my immediate problem, this is a very low priority issue.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> 

I have found that the traced jit module will got dismatched result compared to nnmodule at difference batch inputs as `window_reverse` is not supporting dynamic batch [modeling_swin.py#L218](https://github.com/huggingface/transformers/blob/main/src/transformers/models/swin/modeling_swin.py#L218)：

```python
def window_reverse(windows, window_size, height, width):
    """
    Merges windows to produce higher resolution features.
    """
    batch_size = math.floor(windows.shape[0] / (height * width / window_size / window_size))
    windows = windows.view(batch_size, height // window_size, width // window_size, window_size, window_size, -1)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(batch_size, height, width, -1)
    return windows
```

Is that a possible reason? In my practice，I write `window_reverse` as following for supporting dynamic batch：

```python
def window_reverse(windows, window_size, height, width):
    """
    Merges windows to produce higher resolution features.
    """
    channels = int(windows.shape[-1])
    windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, channels)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, channels)
    return windows
```