Swin Transformer

This project aims to explore the deployment of SwinTransformer based on TensorRT, including the test results of FP16 and INT8.

Introduction(Quoted from the Original Project )

Swin Transformer original github repo (the name Swin stands for Shifted window) is initially described in arxiv, which capably serves as a general-purpose backbone for computer vision. It is basically a hierarchical Transformer whose representation is computed with shifted windows. The shifted windowing scheme brings greater efficiency by limiting self-attention computation to non-overlapping local windows while also allowing for cross-window connection.

Setup

1. Please refer to the Data preparation session to prepare Imagenet-1K.

2. Docker setup.

   a). Docker pull and launch. TensorRT 8.5.1.7 is preinstalled in this docker.

   docker pull nvcr.io/nvidia/tensorrt:22.11-py3
   docker run --name tensorrt_22.11_py3_swin -it --rm --gpus "device=0" --network host --shm-size 16g -v /($path_of_your_projects):/root/space/projects nvcr.io/nvidia/tensorrt:22.11-py3 &
   

   b). Install necessary utils:

   pip install pytorch-quantization==2.1.2 --extra-index-url https://pypi.ngc.nvidia.com
   pip install torch==1.13.0 torchvision==0.14.0
   pip install timm==0.4.12
   pip install termcolor==1.1.0
   pip install pyyaml tqdm yacs
   pip install onnx onnxruntime
   

Code Structure

Focus on the modifications and additions.

.
├── config.py                  # Add the default config of quantization and onnx export
├── export.py                  # Export the PyTorch model to ONNX format
├── calib.sh                   # Calib script
├── models
│   ├── build.py
│   ├── __init__.py
│   └── swin_transformer.py    # Build the model and add the quantization operations, modified to export the onnx and build the TensorRT engine
├── README.md
├── qat.sh                     # Execute calibration and QAT finetuning
├── trt                        # Directory for TensorRT's engine evaluation and visualization.
│   ├── debug                  # Compare scripts with polygraphy, compare the results of onnx and TRT engine with fixed input
│   ├── build_engine.py        # Script for engine build
│   ├── engine.py
│   ├── eval_trt.py            # Evaluate the tensorRT engine's accuary.
│   ├── eval_onnxrt.py         # Run the onnx model, generate the results, just for debugging
├── swin_quant_flow.py         # QAT workflow for swin_transformer, we haven't try the swin_mlp structure
└── weights

Export to ONNX and Build TensorRT Engine

You need to pay attention to some small modifications below.

1. For dynamic batchsize support, please refer to the modifications in models/swin_transformer.py. The window_reverse does not support dynamic batch because it cast the first dimension of windows to integer.

    def window_reverse(windows, window_size, H, W):
        # B = int(windows.shape[0] / (H * W / window_size / window_size))
        # x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
        # x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
        C = int(windows.shape[-1])
        x = windows.view(-1, H // window_size, W // window_size, window_size, window_size, C)
        x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, H, W, C)
        return x
   

2. For fp16 mode, fp16 can't store very large and very small numbers like fp32. So we need to set some specific layers to fp32 during the engine build. We can fallback the POW and REDUCE layers to fp32, it is enough to fix the accuracy problem and don't hurt the perfomance/throughput. Sometime maybe with different weights, you need to fallback POW, REDUCEMEAN, Add and Sqrt to fp16, please refer to fix_fp16_network function in trt/trt_utils.py.
   

Possible issues for old onnx and TensorRT versions

If you are using the env setting as above, just skip this.

1. Exporting the operator roll to ONNX opset version 9 is not supported.
   A: Please refer to torch/onnx/symbolic_opset9.py, add the support of exporting torch.roll.

2. Node (Concat_264) Op (Concat) [ShapeInferenceError] All inputs to Concat must have same rank.
   A: Please refer to the modifications in models/swin_transformer.py. We use the input_resolution and window_size to compute the nW.

      if mask is not None:
        nW = int(self.input_resolution[0]*self.input_resolution[1]/self.window_size[0]/self.window_size[1])
        #nW = mask.shape[0]
        #print('nW: ', nW)
        attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
        attn = attn.view(-1, self.num_heads, N, N)
        attn = self.softmax(attn)
   

Accuray Test Results on ImageNet-1K Validation Dataset

1. Download the Swin-T pretrained model from Model Zoo.

2. export.py exports a pytorch model to onnx format.

   $ python export.py --eval --cfg SwinTransformer/configs/swin/swin_tiny_patch4_window7_224.yaml --resume ./weights/swin_tiny_patch4_window7_224.pth --data-path /root/space/projects/datasets/imagenet  
   

3. Build the TensorRT engine using trtexec.

   $ python trt/build_engine.py --onnx-file ./weights/swin_tiny_patch4_window7_224.onnx --trt-engine  ./weights/swin_tiny_patch4_window7_224_batch32_fp32.engine --verbose --mode fp32 --b-opt 32
   

   For fp16 mode.

   $ python trt/build_engine.py --onnx-file ./weights/swin_tiny_patch4_window7_224.onnx --trt-engine  ./weights/swin_tiny_patch4_window7_224_batch32_fp16.engine --verbose --mode fp16 --b-opt 32
   

   You can use the trtexec to test the throughput of the TensorRT engine.

   $ trtexec --loadEngine=./weights/swin_tiny_patch4_window7_224_batch32_fp16.engine
   

4. trt/eval_trt.py aims to evalute the accuracy of the TensorRT engine.

   $ python trt/eval_trt.py --engine ./weights/swin_tiny_patch4_window7_224_batch32_fp16.engine --data-path /root/space/projects/datasets/imagenet --batch-size 32
   

5. trt/eval_onnxrt.py aims to evalute the accuracy of the Onnx model, just for debug.

   $ python trt/eval_onnxrt.py --eval --cfg SwinTransformer/configs/swin/swin_tiny_patch4_window7_224.yaml --resume ./weights/swin_tiny_patch4_window7_224_fixed.onnx --data-path /root/space/projects/datasets/imagenet --batch-size 32
   

New Test Attached

Accuracy and Speedup Test of TensorRT engine (A100 40GB, TensorRT 8.5.1.7)

| Model (BS=32) | FP32(latency: mean) | FP16(latency: mean) | FP16 Acc.1 | | :---: | :---: | :---: | :---: | Swin Tiny | 14.2938 ms |8.87109 ms |81.20% | Swin Small | 22.9841 ms |12.9888 ms |83.20% | Swin Base | 31.4782 ms |17.3515 ms |85.20% | Swin Large | 53.5593 ms | 27.6452 ms |86.20% |

For int8, after calibration, the accuray is 80.8% with swin-tiny, just as expected. But the speedup is not obvious. So fp16 deployment is highly recommended.

Previous Test Attached

Accuracy Test of TensorRT engine (T4, TensorRT 8.2)

| SwinTransformer(T4) | Acc@1 | Notes | | :---: | :---: | :---: | | PyTorch Pretrained Model | 81.160 | | | TensorRT Engine(FP32) | 81.156 | | | TensorRT Engine(FP16) | 81.150 | With POW and REDUCE layers fallback to FP32 | | TensorRT Engine(INT8 QAT) | - | Finetune for 1 epoch, got 79.980, need to improve the int8 throughput first |

Speed Test of TensorRT engine (T4, TensorRT 8.2)

| SwinTransformer(T4) | FP32 | FP16 | Explicit Quantization(INT8, QAT) | | :---: | :---: | :---: | :---: | | batchsize=1 | 245.388 qps | 510.072 qps | 385.454 qps | | batchsize=16 | 316.8624 qps | 804.112 qps | 815.606 qps | | batchsize=64 | 329.13984 qps | 833.4208 qps | 780.006 qps | | batchsize=256 | 331.9808 qps | 844.10752 qps | - |

Result:

1. Now the accuracy and speedup of FP16 is as expected, it is highly recommended to deploy Swin-Transformer with FP16 precision.

2. Compared with FP16, INT8 does not speed up at present.

Add Quantizer and Wrap the Fake-Quantized Model (Experiment)

The main modifications of models/swin_transformer.py are as below.

1. For PatchMerging block, modify torch.nn.Liner to quant_nn.QuantLinear.

2. For WindowAttention block,
   a) For query, key and value, modify torch.nn.Liner to quant_nn.QuantLinear.
   b) Quantize the four inputs of torch.matmul.

3. For MLP block, modify torch.nn.Liner to quant_nn.QuantLinear.

4. For SwinTransformerBlock block, quantize the inputs of operator +.

QAT for Swin Transformer (Experiment)

In order to do the QAT finetuning, some utils are needed to install.
tqdm, prettytable, scipy, absl-py

1. With swin_quant_flow.py, wrap a fake-quantized model, calibrate, QAT finetuning and export to onnx model.

   $ ./calib.sh
   

   Or you can run calibration and QAT-finetuning in the same time.

   $ ./qat.sh
   

2. Build TensorRT engine and evaluate as above. Same commands.