Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/cars.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/cell.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/demo_3x2.gif filter=lfs diff=lfs merge=lfs -text
+assets/top.gif filter=lfs diff=lfs merge=lfs -text

aot/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020, z-x-yang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

aot/MODEL_ZOO.md

@@ -0,0 +1,115 @@
+## Model Zoo and Results
+
+### Environment and Settings
+- 4/1 NVIDIA V100 GPUs for training/evaluation.
+- Auto-mixed precision was enabled in training but disabled in evaluation.
+- Test-time augmentations were not used.
+- The inference resolution of DAVIS/YouTube-VOS was 480p/1.3x480p as [CFBI](https://github.com/z-x-yang/CFBI).
+- Fully online inference. We passed all the modules frame by frame.
+- Multi-object FPS was recorded instead of single-object one.
+
+### Pre-trained Models
+Stages:
+
+- `PRE`: the pre-training stage with static images.
+
+- `PRE_YTB_DAV`: the main-training stage with YouTube-VOS and DAVIS. All the kinds of evaluation share an **identical** model and the **same** parameters.
+
+
+| Model      | Param (M) |                                             PRE                                              |                                         PRE_YTB_DAV                                          |
+|:---------- |:---------:|:--------------------------------------------------------------------------------------------:|:--------------------------------------------------------------------------------------------:|
+| AOTT       |    5.7    | [gdrive](https://drive.google.com/file/d/1_513h8Hok9ySQPMs_dHgX5sPexUhyCmy/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1owPmwV4owd_ll6GuilzklqTyAd0ZvbCu/view?usp=sharing) |
+| AOTS       |    7.0    | [gdrive](https://drive.google.com/file/d/1QUP0-VED-lOF1oX_ppYWnXyBjvUzJJB7/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1beU5E6Mdnr_pPrgjWvdWurKAIwJSz1xf/view?usp=sharing) |
+| AOTB       |    8.3    | [gdrive](https://drive.google.com/file/d/11Bx8n_INAha1IdpHjueGpf7BrKmCJDvK/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1hH-GOn4GAxHkV8ARcQzsUy8Ax6ndot-A/view?usp=sharing) |
+| AOTL       |    8.3    | [gdrive](https://drive.google.com/file/d/1WL6QCsYeT7Bt-Gain9ZIrNNXpR2Hgh29/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1L1N2hkSPqrwGgnW9GyFHuG59_EYYfTG4/view?usp=sharing) |
+| R50-AOTL   |   14.9    | [gdrive](https://drive.google.com/file/d/1hS4JIvOXeqvbs-CokwV6PwZV-EvzE6x8/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1qJDYn3Ibpquu4ffYoQmVjg1YCbr2JQep/view?usp=sharing) |
+| SwinB-AOTL |   65.4    | [gdrive](https://drive.google.com/file/d/1LlhKQiXD8JyZGGs3hZiNzcaCLqyvL9tj/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/192jCGQZdnuTsvX-CVra-KVZl2q1ZR0vW/view?usp=sharing) |
+
+| Model      | Param (M) |                                             PRE                                              |                                         PRE_YTB_DAV                                          |
+|:---------- |:---------:|:--------------------------------------------------------------------------------------------:|:--------------------------------------------------------------------------------------------:|
+| DeAOTT       |    7.2   | [gdrive](https://drive.google.com/file/d/11C1ZBoFpL3ztKtINS8qqwPSldfYXexFK/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1ThWIZQS03cYWx1EKNN8MIMnJS5eRowzr/view?usp=sharing) |
+| DeAOTS       |    10.2   | [gdrive](https://drive.google.com/file/d/1uUidrWVoaP9A5B5-EzQLbielUnRLRF3j/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1YwIAV5tBtn5spSFxKLBQBEQGwPHyQlHi/view?usp=sharing) |
+| DeAOTB       |    13.2   | [gdrive](https://drive.google.com/file/d/1bEQr6vIgQMVITrSOtxWTMgycKpS0cor9/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1BHxsonnvJXylqHlZ1zJHHc-ymKyq-CFf/view?usp=sharing) |
+| DeAOTL       |    13.2   | [gdrive](https://drive.google.com/file/d/1_vBL4KJlmBy0oBE4YFDOvsYL1ZtpEL32/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/18elNz_wi9JyVBcIUYKhRdL08MA-FqHD5/view?usp=sharing) |
+| R50-DeAOTL   |    19.8   | [gdrive](https://drive.google.com/file/d/1sTRQ1g0WCpqVCdavv7uJiZNkXunBt3-R/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1QoChMkTVxdYZ_eBlZhK2acq9KMQZccPJ/view?usp=sharing) |
+| SwinB-DeAOTL |    70.3   | [gdrive](https://drive.google.com/file/d/16BZEE53no8CxT-pPLDC2q1d6Xlg8mWPU/view?usp=sharing) | [gdrive](https://drive.google.com/file/d/1g4E-F0RPOx9Nd6J7tU9AE1TjsouL4oZq/view?usp=sharing) |
+
+To use our pre-trained model to infer, a simple way is to set `--model` and `--ckpt_path` to your downloaded checkpoint's model type and file path when running `eval.py`.
+
+### YouTube-VOS 2018 val
+`ALL-F`: all frames. The default evaluation setting of YouTube-VOS is 6fps, but 30fps sequences (all the frames) are also supplied by the dataset organizers. We noticed that many VOS methods prefer to evaluate with 30fps videos. Thus, we also supply our results here. Denser video sequences can significantly improve VOS performance when using the memory reading strategy (like AOTL, R50-AOTL, and SwinB-AOTL), but the efficiency will be influenced since more memorized frames are stored for object matching.
+| Model        |    Stage    |   FPS    | All-F |   Mean   |  J Seen  |  F Seen  | J Unseen | F Unseen |                                         Predictions                                          |
+|:------------ |:-----------:|:--------:|:-----:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------------------------------------------------------------------------------------------:|
+| AOTT         | PRE_YTB_DAV |   41.0   |       |   80.2   |   80.4   |   85.0   |   73.6   |   81.7   | [gdrive](https://drive.google.com/file/d/1u8mvPRT08ENZHsw9Xf_4C6Sv9BoCzENR/view?usp=sharing) |
+| AOTT         | PRE_YTB_DAV |   41.0   |   √   |   80.9   |   80.0   |   84.7   |   75.2   |   83.5   | [gdrive](https://drive.google.com/file/d/1RGMI5-29Z0odq73rt26eCxOUYUd-fvVv/view?usp=sharing) |
+| DeAOTT       | PRE_YTB_DAV | **53.4** |       | **82.0** | **81.6** | **86.3** | **75.8** | **84.2** |    -                                                                                          |
+| AOTS         | PRE_YTB_DAV |   27.1   |       |   82.9   |   82.3   |   87.0   |   77.1   |   85.1   | [gdrive](https://drive.google.com/file/d/1a4-rNnxjMuPBq21IKo31WDYZXMPgS7r2/view?usp=sharing) |
+| AOTS         | PRE_YTB_DAV |   27.1   |   √   |   83.0   |   82.2   |   87.0   |   77.3   |   85.7   | [gdrive](https://drive.google.com/file/d/1Z0cndyoCw5Na6u-VFRE8CyiIG2RbMIUO/view?usp=sharing) |
+| DeAOTS       | PRE_YTB_DAV | **38.7** |       | **84.0** | **83.3** | **88.3** | **77.9** | **86.6** |    -                                                                                          |
+| AOTB         | PRE_YTB_DAV |   20.5   |       |   84.0   |   83.2   |   88.1   |   78.0   |   86.5   | [gdrive](https://drive.google.com/file/d/1J5nhuQbbjVLYNXViBIgo21ddQy-MiOLG/view?usp=sharing) |
+| AOTB         | PRE_YTB_DAV |   20.5   |   √   |   84.1   |   83.6   |   88.5   |   78.0   |   86.5   | [gdrive](https://drive.google.com/file/d/1gFaweB_GTJjHzSD61v_ZsY9K7UEND30O/view?usp=sharing) |
+| DeAOTB       | PRE_YTB_DAV | **30.4** |       | **84.6** | **83.9** | **88.9** | **78.5** | **87.0** |    -                                                                                          |
+| AOTL         | PRE_YTB_DAV |   16.0   |       |   84.1   |   83.2   |   88.2   |   78.2   |   86.8   | [gdrive](https://drive.google.com/file/d/1kS8KWQ2L3wzxt44ROLTxwZOT7ZpT8Igc/view?usp=sharing) |
+| AOTL         | PRE_YTB_DAV |   6.5    |   √   |   84.5   |   83.7   |   88.8   |   78.4   | **87.1** | [gdrive](https://drive.google.com/file/d/1Rpm3e215kJOUvb562lJ2kYg2I3hkrxiM/view?usp=sharing) |
+| DeAOTL       | PRE_YTB_DAV | **24.7** |       | **84.8** | **84.2** | **89.4** | **78.6** |   87.0   |    -                                                                                          |
+| R50-AOTL     | PRE_YTB_DAV |   14.9   |       |   84.6   |   83.7   |   88.5   |   78.8   |   87.3   | [gdrive](https://drive.google.com/file/d/1nbJZ1bbmEgyK-bg6HQ8LwCz5gVJ6wzIZ/view?usp=sharing) |
+| R50-AOTL     | PRE_YTB_DAV |   6.4    |   √   |   85.5   |   84.5   |   89.5   |   79.6   |   88.2   | [gdrive](https://drive.google.com/file/d/1NbB54ZhYvfJh38KFOgovYYPjWopd-2TE/view?usp=sharing) |
+| R50-DeAOTL   | PRE_YTB_DAV | **22.4** |       | **86.0** | **84.9** | **89.9** | **80.4** | **88.7** |    -                                                                                         |
+| SwinB-AOTL   | PRE_YTB_DAV |   9.3    |       |   84.7   |   84.5   |   89.5   |   78.1   |   86.7   | [gdrive](https://drive.google.com/file/d/1QFowulSY0LHfpsjUV8ZE9rYc55L9DOC7/view?usp=sharing) |
+| SwinB-AOTL   | PRE_YTB_DAV |   5.2    |   √   |   85.1   |   85.1   |   90.1   |   78.4   |   86.9   | [gdrive](https://drive.google.com/file/d/1TulhVOhh01rkssNYbOQASeWKu7CQ5Azx/view?usp=sharing) |
+| SwinB-DeAOTL | PRE_YTB_DAV | **11.9** |       | **86.2** | **85.6** | **90.6** | **80.0** | **88.4** |    -                                                                                  |
+
+### YouTube-VOS 2019 val
+| Model        |    Stage    |   FPS    | All-F |   Mean   |  J Seen  |  F Seen  | J Unseen | F Unseen |                                         Predictions                                          |
+|:------------ |:-----------:|:--------:|:-----:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------------------------------------------------------------------------------------------:|
+| AOTT         | PRE_YTB_DAV |   41.0   |       |   80.0   |   79.8   |   84.2   |   74.1   |   82.1   | [gdrive](https://drive.google.com/file/d/1zzyhN1XYtajte5nbZ7opOdfXeDJgCxC5/view?usp=sharing) |
+| AOTT         | PRE_YTB_DAV |   41.0   |   √   |   80.9   |   79.9   |   84.4   |   75.6   |   83.8   | [gdrive](https://drive.google.com/file/d/1V_5vi9dAXOis_WrDieacSESm7OX20Bv-/view?usp=sharing) |
+| DeAOTT       | PRE_YTB_DAV | **53.4** |       | **82.0** | **81.2** | **85.6** | **76.4** | **84.7** |     -                                                                                         |
+| AOTS         | PRE_YTB_DAV |   27.1   |       |   82.7   |   81.9   |   86.5   |   77.3   |   85.2   | [gdrive](https://drive.google.com/file/d/11YdkUeyjkTv8Uw7xMgPCBzJs6v5SDt6n/view?usp=sharing) |
+| AOTS         | PRE_YTB_DAV |   27.1   |   √   |   82.8   |   81.9   |   86.5   |   77.3   |   85.6   | [gdrive](https://drive.google.com/file/d/1UhyurGTJeAw412czU3_ebzNwF8xQ4QG_/view?usp=sharing) |
+| DeAOTS       | PRE_YTB_DAV | **38.7** |       | **83.8** | **82.8** | **87.5** | **78.1** | **86.8** |     -                                                                                         |
+| AOTB         | PRE_YTB_DAV |   20.5   |       |   84.0   |   83.1   |   87.7   |   78.5   |   86.8   | [gdrive](https://drive.google.com/file/d/1NeI8cT4kVqTqVWAwtwiga1rkrvksNWaO/view?usp=sharing) |
+| AOTB         | PRE_YTB_DAV |   20.5   |   √   |   84.1   |   83.3   |   88.0   |   78.2   |   86.7   | [gdrive](https://drive.google.com/file/d/1kpYV2XFR0sOfLWD-wMhd-nUO6CFiLjlL/view?usp=sharing) |
+| DeAOTB       | PRE_YTB_DAV | **30.4** |       | **84.6** | **83.5** | **88.3** | **79.1** | **87.5** |     -                                                                                         |
+| AOTL         | PRE_YTB_DAV |   16.0   |       |   84.0   |   82.8   |   87.6   |   78.6   |   87.1   | [gdrive](https://drive.google.com/file/d/1qKLlNXxmT31bW0weEHI_zAf4QwU8Lhou/view?usp=sharing) |
+| AOTL         | PRE_YTB_DAV |   6.5    |   √   |   84.2   |   83.0   |   87.8   |   78.7   |   87.3   | [gdrive](https://drive.google.com/file/d/1o3fwZ0cH71bqHSA3bYNjhP4GGv9Vyuwa/view?usp=sharing) |
+| DeAOTL       | PRE_YTB_DAV | **24.7** |       | **84.7** | **83.8** | **88.8** | **79.0** | **87.2** |     -                                                                                         |
+| R50-AOTL     | PRE_YTB_DAV |   14.9   |       |   84.4   |   83.4   |   88.1   |   78.7   |   87.2   | [gdrive](https://drive.google.com/file/d/1I7ooSp8EYfU6fvkP6QcCMaxeencA68AH/view?usp=sharing) |
+| R50-AOTL     | PRE_YTB_DAV |   6.4    |   √   |   85.3   |   83.9   |   88.8   |   79.9   |   88.5   | [gdrive](https://drive.google.com/file/d/1OGqlkEu0uXa8QVWIVz_M5pmXXiYR2sh3/view?usp=sharing) |
+| R50-DeAOTL   | PRE_YTB_DAV | **22.4** |       | **85.9** | **84.6** | **89.4** | **80.8** | **88.9** |     -                                                                                         |
+| SwinB-AOTL   | PRE_YTB_DAV |   9.3    |       |   84.7   |   84.0   |   88.8   |   78.7   |   87.1   | [gdrive](https://drive.google.com/file/d/1fPzCxi5GM7N2sLKkhoTC2yoY_oTQCHp1/view?usp=sharing) |
+| SwinB-AOTL   | PRE_YTB_DAV |   5.2    |   √   |   85.3   |   84.6   |   89.5   |   79.3   |   87.7   | [gdrive](https://drive.google.com/file/d/1e3D22s_rJ7Y2X2MHo7x5lcNtwmHFlwYB/view?usp=sharing) |
+| SwinB-DeAOTL | PRE_YTB_DAV | **11.9** |       | **86.1** | **85.3** | **90.2** | **80.4** | **88.6** |     -                                                                                         |
+
+### DAVIS-2017 test
+
+| Model      |    Stage    | FPS  |   Mean   | J Score  | F Score  | Predictions |
+| ---------- |:-----------:|:----:|:--------:|:--------:|:--------:|:----:|
+| AOTT       | PRE_YTB_DAV | **51.4** |   73.7   |   70.0   |   77.3   | [gdrive](https://drive.google.com/file/d/14Pu-6Uz4rfmJ_WyL2yl57KTx_pSSUNAf/view?usp=sharing) |
+| AOTS       | PRE_YTB_DAV | 40.0 |   75.2   |   71.4   |   78.9   | [gdrive](https://drive.google.com/file/d/1zzAPZCRLgnBWuAXqejPPEYLqBxu67Rj1/view?usp=sharing) |
+| AOTB       | PRE_YTB_DAV | 29.6 |   77.4   |   73.7   |   81.1   | [gdrive](https://drive.google.com/file/d/1WpQ-_Jrs7Ssfw0oekrejM2OVWEx_tBN1/view?usp=sharing) |
+| AOTL       | PRE_YTB_DAV | 18.7 |   79.3   |   75.5   |   83.2   | [gdrive](https://drive.google.com/file/d/1rP1Zdgc0N1d8RR2EaXMz3F-o5zqcNVe8/view?usp=sharing) |
+| R50-AOTL   | PRE_YTB_DAV | 18.0 |   79.5   |   76.0   |   83.0   | [gdrive](https://drive.google.com/file/d/1iQ5iNlvlS-In586ZNc4LIZMSdNIWDvle/view?usp=sharing) |
+| SwinB-AOTL | PRE_YTB_DAV | 12.1  | **82.1** | **78.2** | **85.9** | [gdrive](https://drive.google.com/file/d/1oVt4FPcZdfVHiOxjYYKef0q7Ovy4f5Q_/view?usp=sharing) |
+
+### DAVIS-2017 val
+
+| Model      |    Stage    | FPS  |   Mean   | J Score  |  F Score  | Predictions |
+| ---------- |:-----------:|:----:|:--------:|:--------:|:---------:|:----:|
+| AOTT       | PRE_YTB_DAV | **51.4** |   79.2   |   76.5   |   81.9    | [gdrive](https://drive.google.com/file/d/10OUFhK2Sz-hOJrTDoTI0mA45KO1qodZt/view?usp=sharing) |
+| AOTS       | PRE_YTB_DAV | 40.0 |   82.1   |   79.3   |   84.8    | [gdrive](https://drive.google.com/file/d/1T-JTYyksWlq45jxcLjnRaBvvYUhWgHFH/view?usp=sharing) |
+| AOTB       | PRE_YTB_DAV | 29.6 |   83.3   |   80.6   |   85.9    | [gdrive](https://drive.google.com/file/d/1EVUnxQm9TLBTuwK82QyiSKk9R9V8NwRL/view?usp=sharing) |
+| AOTL       | PRE_YTB_DAV | 18.7 |   83.6   |   80.8   |   86.3    | [gdrive](https://drive.google.com/file/d/1CFauSni2BxAe_fcl8W_6bFByuwJRbDYm/view?usp=sharing) |
+| R50-AOTL   | PRE_YTB_DAV | 18.0 |   85.2   |   82.5   |   87.9    | [gdrive](https://drive.google.com/file/d/1vjloxnP8R4PZdsH2DDizfU2CrkdRHHyo/view?usp=sharing) |
+| SwinB-AOTL | PRE_YTB_DAV | 12.1  | **85.9** | **82.9** | **88.9** | [gdrive](https://drive.google.com/file/d/1tYCbKOas0i7Et2iyUAyDwaXnaD9YWxLr/view?usp=sharing) |
+
+### DAVIS-2016 val
+
+| Model      |    Stage    | FPS  |   Mean   | J Score  | F Score  | Predictions |
+| ---------- |:-----------:|:----:|:--------:|:--------:|:--------:|:----:|
+| AOTT       | PRE_YTB_DAV | **51.4** |   87.5   |   86.5   |   88.4   | [gdrive](https://drive.google.com/file/d/1LeW8WQhnylZ3umT7E379KdII92uUsGA9/view?usp=sharing) |
+| AOTS       | PRE_YTB_DAV | 40.0 |   89.6   |   88.6   |   90.5   | [gdrive](https://drive.google.com/file/d/1vqGei5tLu1FPVrTi5bwRAsaGy3Upf7B1/view?usp=sharing) |
+| AOTB       | PRE_YTB_DAV | 29.6 |   90.9   |   89.6   |   92.1   | [gdrive](https://drive.google.com/file/d/1qAppo2uOVu0FbE9t1FBUpymC3yWgw1LM/view?usp=sharing) |
+| AOTL       | PRE_YTB_DAV | 18.7 |   91.1   |   89.5   |   92.7   | [gdrive](https://drive.google.com/file/d/1g6cjYhgBWjMaY3RGAm31qm3SPEF3QcKV/view?usp=sharing) |
+| R50-AOTL   | PRE_YTB_DAV | 18.0 |   91.7   |   90.4   |   93.0   | [gdrive](https://drive.google.com/file/d/1QzxojqWKsvRf53K2AgKsK523ZVuYU4O-/view?usp=sharing) |
+| SwinB-AOTL | PRE_YTB_DAV | 12.1  | **92.2** | **90.6** | **93.8** | [gdrive](https://drive.google.com/file/d/1RIqUtAyVnopeogfT520d7a0yiULg1obp/view?usp=sharing) |

aot/Pytorch-Correlation-extension/.gitignore

@@ -0,0 +1 @@
+*.egg*

aot/Pytorch-Correlation-extension/Correlation_Module/correlation.cpp

@@ -0,0 +1,178 @@
+#include <torch/extension.h>
+using namespace torch;
+
+#include <vector>
+
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+template <typename scalar_t>
+static void correlate_patch(
+    TensorAccessor<scalar_t,3> input1,
+    TensorAccessor<scalar_t,3> input2,
+    scalar_t *dst,
+    int kH, int kW,
+    int dilationH, int dilationW,
+    int u, int v,
+    int shiftU, int shiftV){
+  const int C = input1.size(0);
+  const int iH = input1.size(1);
+  const int iW = input1.size(2);
+  for (int c=0; c<C; ++c){
+    for (int i=0; i<kH; ++i){
+      int i1 = u + i * dilationH;
+      int i2 = i1 + shiftU;
+      if WITHIN_BOUNDS(i1, i2, iH, iH){
+        for (int j=0; j<kW; ++j){
+          int j1 = v + j * dilationW;
+          int j2 = j1 + shiftV;
+          if WITHIN_BOUNDS(j1, j2, iW, iW){
+            scalar_t v1 = input1[c][i1][j1];
+            scalar_t v2 = input2[c][i2][j2];
+            *dst += v1 * v2;
+          }
+        }
+      }
+    }
+  }
+}
+
+template <typename scalar_t>
+static void correlate_patch_grad(
+    TensorAccessor<scalar_t,3> input1,
+    TensorAccessor<scalar_t,3> gradInput1,
+    TensorAccessor<scalar_t,3> input2,
+    TensorAccessor<scalar_t,3> gradInput2,
+    scalar_t gradOutput,
+    int kH, int kW,
+    int dilationH, int dilationW,
+    int u, int v,
+    int shiftU, int shiftV){
+
+  const int C = input1.size(0);
+  const int iH = input1.size(1);
+  const int iW = input1.size(2);
+
+  for (int c=0; c<C; ++c){
+    for (int i=0; i<kH; ++i){
+      int i1 = u + i * dilationH;
+      int i2 = i1 + shiftU;
+      if WITHIN_BOUNDS(i1, i2, iH, iH){
+        for (int j=0; j<kW; ++j){
+          int j1 = v + j * dilationW;
+          int j2 = j1 + shiftV;
+          if WITHIN_BOUNDS(j1, j2, iW, iW){
+            scalar_t v1 = input1[c][i1][j1];
+            scalar_t v2 = input2[c][i2][j2];
+            gradInput2[c][i2][j2] += gradOutput * v1;
+            gradInput1[c][i1][j1] += gradOutput * v2;
+          }
+        }
+      }
+    }
+  }
+}
+
+torch::Tensor correlation_cpp_forward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+
+  const auto batch_size = input1.size(0);
+  const auto iH = input1.size(2);
+  const auto iW = input1.size(3);
+  const int patchRadH = (patchH - 1) / 2;
+  const int patchRadW = (patchW - 1) / 2;
+  const int dilatedKH = (kH - 1) * dilationH + 1;
+  const int dilatedKW = (kW - 1) * dilationW + 1;
+
+  const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+  const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+  auto output = at::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+
+  int n, ph, pw, h, w;
+  #pragma omp parallel for private(n, ph, pw, h, w) collapse(2)
+    for (n = 0; n < batch_size; ++n) {
+      for(ph = 0; ph < patchH; ++ph){
+        for(pw = 0; pw < patchW; ++pw){
+          AT_DISPATCH_FLOATING_TYPES(input1.scalar_type(), "correlation_forward_cpp", ([&] {
+            auto input1_acc = input1.accessor<scalar_t, 4>();
+            auto input2_acc = input2.accessor<scalar_t, 4>();
+            auto output_acc = output.accessor<scalar_t, 5>();
+            for (h = 0; h < oH; ++h) {
+              for (w = 0; w < oW; ++w) {
+                correlate_patch(input1_acc[n],
+                                input2_acc[n],
+                                &output_acc[n][ph][pw][h][w],
+                                kH, kW,
+                                dilationH, dilationW,
+                                -padH + h * dH,
+                                -padW + w * dW,
+                                (ph - patchRadH)  * dilation_patchH,
+                                (pw - patchRadW)  * dilation_patchW);
+              }
+            }
+          }));
+        }
+      }
+    }
+  return output;
+}
+
+std::vector<torch::Tensor> correlation_cpp_backward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    torch::Tensor gradOutput,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+  
+  const int batch_size = input1.size(0);
+  const int patchRadH = (patchH - 1) / 2;
+  const int patchRadW = (patchW - 1) / 2;
+  const int oH = gradOutput.size(3);
+  const int oW = gradOutput.size(4);
+  
+  auto gradInput1 = torch::zeros_like(input1);
+
+  auto gradInput2 = torch::zeros_like(input2);
+
+  int n, ph, pw, h, w;
+  #pragma omp parallel for private(n, ph, pw, h, w)
+    for (n = 0; n < batch_size; ++n) {
+      AT_DISPATCH_FLOATING_TYPES(input1.scalar_type(), "correlation_backward_cpp", ([&] {
+        auto input1_acc = input1.accessor<scalar_t, 4>();
+        auto gradInput1_acc = gradInput1.accessor<scalar_t, 4>();
+        auto input2_acc = input2.accessor<scalar_t, 4>();
+        auto gradInput2_acc = gradInput2.accessor<scalar_t, 4>();
+        auto gradOutput_acc = gradOutput.accessor<scalar_t, 5>();
+
+        for(ph = 0; ph < patchH; ++ph){
+          for(pw = 0; pw < patchW; ++pw){
+            for (h = 0; h < oH; ++h) {
+              for (w = 0; w < oW; ++w) {
+                correlate_patch_grad(input1_acc[n], gradInput1_acc[n],
+                                     input2_acc[n], gradInput2_acc[n],
+                                     gradOutput_acc[n][ph][pw][h][w],
+                                     kH, kW,
+                                     dilationH, dilationW,
+                                     -padH + h * dH,
+                                     -padW + w * dW,
+                                     (ph - patchRadH)  * dilation_patchH,
+                                     (pw - patchRadW)  * dilation_patchW);
+              }
+            }
+          }
+        }
+      }));
+    }
+
+  return {gradInput1, gradInput2};
+}

aot/Pytorch-Correlation-extension/Correlation_Module/correlation_cuda_kernel.cu

@@ -0,0 +1,327 @@
+#include <torch/types.h>
+using namespace torch;
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+#include <iostream>
+
+// Cuda tensor accessor definitions
+// restrict pointer traits piroritize speed over memory consumption
+#define TensorAcc4R PackedTensorAccessor32<scalar_t,4,RestrictPtrTraits>
+#define TensorAcc5R PackedTensorAccessor32<scalar_t,5,RestrictPtrTraits>
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+#define THREADS_FORWARD 32
+#define THREADS_BACKWARD 5
+
+
+namespace corr {
+template <typename scalar_t>
+__global__ void correlation_cuda_forward_kernel(
+    const TensorAcc4R rInput1,
+    const TensorAcc4R rInput2,
+    TensorAcc5R output,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+  
+  const int iH = rInput1.size(1);
+  const int iW = rInput1.size(2);
+  const int C = rInput1.size(3);
+
+  const int n = blockIdx.x;
+  const int h = blockIdx.y;
+  const int w = blockIdx.z;
+  const int thread = threadIdx.x;
+
+  const int start_i = -padH + h * dH;
+  const int start_j = -padW + w * dW;
+
+  const int patchRadH = dilation_patchH * (patchH - 1) / 2;
+  const int patchRadW = dilation_patchW * (patchW - 1) / 2;
+
+  __shared__ scalar_t prod_sum[THREADS_FORWARD];
+
+  for(int ph = 0; ph < patchH; ++ph){
+    int ph_dilated = ph * dilation_patchH - patchRadH;
+    for(int pw = 0; pw < patchW; ++pw){
+      int pw_dilated = pw * dilation_patchW - patchRadW;
+      prod_sum[thread] = 0;
+      for (int i=0; i<kH; ++i){
+        int i1 = start_i + i * dilationH;
+        int i2 = i1 + ph_dilated;
+        if WITHIN_BOUNDS(i1, i2, iH, iH){
+          for (int j=0; j<kW; ++j){
+            int j1 = start_j + j * dilationW;
+            int j2 = j1 + pw_dilated;
+            if WITHIN_BOUNDS(j1, j2, iW, iW){
+              for (int c=thread; c<C; c += THREADS_FORWARD){
+                scalar_t v1 = rInput1[n][i1][j1][c];
+                scalar_t v2 = rInput2[n][i2][j2][c];
+                prod_sum[thread] += v1 * v2;
+              }
+            }
+          }
+        }
+      }
+      // accumulate 
+      __syncthreads();
+      if (thread == 0) {
+        scalar_t reduce_sum = 0;
+        for (int index = 0; index < THREADS_FORWARD; ++index) {
+          reduce_sum += prod_sum[index];
+        }
+        output[n][ph][pw][h][w] = reduce_sum;
+      }
+    }
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input1(
+    const TensorAcc5R gradOutput,
+    const TensorAcc4R input2,
+    TensorAcc4R gradInput1,
+    const int kH, const int kW,
+    const int patchH, const int patchW,
+    const int padH, const int padW,
+    const int dilationH, const int dilationW,
+    const int dilation_patchH, const int dilation_patchW,
+    const int dH, const int dW,
+    const int batch) {
+  const int iH = input2.size(2);
+  const int iW = input2.size(3);
+
+  const int H = gradOutput.size(3);
+  const int W = gradOutput.size(4);
+
+  const int patchRadH = (patchH - 1) / 2;
+  const int patchRadW = (patchW - 1) / 2;
+  
+  const int n = batch;
+  const int c = blockIdx.x;
+  const int h = blockIdx.y;
+  const int w = blockIdx.z;
+  const int ph_off = threadIdx.x;
+  const int pw_off = threadIdx.y;
+
+  const int h_2 = h + padH;
+  const int w_2 = w + padW;
+  const int min_h = h_2 - kH * dilationH;
+  const int min_w = w_2 - kW * dilationW;
+
+  __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+  prod_sum[ph_off][pw_off] = 0;
+
+  for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+    int i1 = h + dilation_patchH * (ph - patchRadH);
+    for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+      int j1 = w + dilation_patchW * (pw - patchRadW);
+      if (WITHIN_BOUNDS(i1, j1, iH, iW)){
+        scalar_t val = input2[n][c][i1][j1];
+        for(int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+          int i2 = (h_3)/dH;
+          if (i2 * dH != h_3)
+            continue;
+          for(int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+            int j2 = (w_3) / dW;
+            if(j2 * dW != w_3)
+              continue;
+            if WITHIN_BOUNDS(i2, j2, H, W) {
+              prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  __syncthreads();
+
+  if (ph_off == 0 && pw_off == 0){
+    scalar_t reduce_sum =0;
+    for (int ph = 0; ph < THREADS_BACKWARD; ++ph){
+      for (int pw = 0; pw < THREADS_BACKWARD; ++pw){
+        reduce_sum += prod_sum[ph][pw];
+      }
+    }
+    gradInput1[n][c][h][w] = reduce_sum;
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input2(
+    const TensorAcc5R gradOutput,
+    const TensorAcc4R input1,
+    TensorAcc4R gradInput2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW,
+    int batch) {
+  const int iH = input1.size(2);
+  const int iW = input1.size(3);
+
+  const int patchRadH = (patchH - 1) / 2;
+  const int patchRadW = (patchW - 1) / 2;
+
+  const int H = gradOutput.size(3);
+  const int W = gradOutput.size(4);
+
+  const int dilatedKH = kH * dilationH;
+  const int dilatedKW = kW * dilationW;
+  
+  const int n = batch;
+  const int c = blockIdx.x;
+  const int h = blockIdx.y;
+  const int w = blockIdx.z;
+  const int ph_off = threadIdx.x;
+  const int pw_off = threadIdx.y;
+
+  __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+  prod_sum[ph_off][pw_off] = 0;
+
+  for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+    int i1 = h - dilation_patchH * (ph - patchRadH);
+    for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+      int j1 = w - dilation_patchW * (pw - patchRadW);
+      if WITHIN_BOUNDS(i1, j1, iH, iW) {
+        scalar_t val = input1[n][c][i1][j1];
+        
+        const int h_2 = i1 + padH;
+        const int w_2 = j1 + padW;
+        const int min_h = h_2 - dilatedKH;
+        const int min_w = w_2 - dilatedKW;
+        
+        for(int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+          int i2 = (h_3)/dH;
+          if (i2 * dH != h_3)
+            continue;
+          for(int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+            int j2 = (w_3) / dW;
+            if(j2 * dW != w_3)
+              continue;
+            if WITHIN_BOUNDS(i2, j2, H, W) {
+              prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  __syncthreads();
+
+  if (ph_off == 0 && pw_off == 0){
+    scalar_t reduce_sum =0;
+    for (int ph = 0; ph < THREADS_BACKWARD; ++ph){
+      for (int pw = 0; pw < THREADS_BACKWARD; ++pw){
+        reduce_sum += prod_sum[ph][pw];
+      }
+    }
+    gradInput2[n][c][h][w] = reduce_sum;
+  }
+}
+} // namsepace corr
+
+torch::Tensor correlation_cuda_forward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+  
+  const int batch_size = input1.size(0);
+  const int iH = input1.size(2);
+  const int iW = input1.size(3);
+  const int dilatedKH = (kH - 1) * dilationH + 1;
+  const int dilatedKW = (kW - 1) * dilationW + 1;
+
+  const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+  const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+  auto output = torch::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+  
+  auto trInput1 = input1.permute({0, 2, 3, 1}).contiguous();
+  auto trInput2 = input2.permute({0, 2, 3, 1}).contiguous();
+  
+  const int threads = THREADS_FORWARD;
+  const dim3 blocks(batch_size, oH, oW);
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.scalar_type(), "correlation_forward_cuda", ([&] {
+    TensorAcc4R trInput1_acc  = trInput1.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc4R trInput2_acc = trInput2.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc5R output_acc = output.packed_accessor32<scalar_t,5,RestrictPtrTraits>();
+    corr::correlation_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
+        trInput1_acc, trInput2_acc, output_acc,
+        kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
+        dilation_patchH, dilation_patchW, dH, dW);
+  }));
+
+  return output;
+}
+
+std::vector<torch::Tensor> correlation_cuda_backward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    torch::Tensor gradOutput,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+  
+  auto gradInput1 = torch::zeros_like(input1);
+  auto gradInput2 = torch::zeros_like(input2);
+
+  const int batch_size = input1.size(0);
+  const int iH = input1.size(2);
+  const int iW = input1.size(3);
+  const int C = input1.size(1);
+
+  const dim3 blocks(C, iH, iW);
+  const dim3 threads(THREADS_BACKWARD, THREADS_BACKWARD);
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.scalar_type(), "correlation_backward_cuda", ([&] {
+    TensorAcc4R input1_acc = input1.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc4R input2_acc = input2.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc4R gradInput1_acc = gradInput1.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc4R gradInput2_acc = gradInput2.packed_accessor32<scalar_t,4,RestrictPtrTraits>();
+    TensorAcc5R gradOutput_acc = gradOutput.packed_accessor32<scalar_t,5,RestrictPtrTraits>();
+
+
+    for (int n = 0; n < batch_size; ++n){
+      corr::correlation_cuda_backward_kernel_input1<scalar_t><<<blocks, threads>>>(
+          gradOutput_acc, input2_acc, gradInput1_acc,
+          kH, kW, patchH, patchW, padH, padW,
+          dilationH, dilationW,
+          dilation_patchH, dilation_patchW, 
+          dH, dW,
+          n);
+    }
+
+    for (int n = 0; n < batch_size; ++n){
+      corr::correlation_cuda_backward_kernel_input2<scalar_t><<<blocks, threads>>>(
+          gradOutput_acc, input1_acc, gradInput2_acc,
+          kH, kW, patchH, patchW, padH, padW,
+          dilationH, dilationW,
+          dilation_patchH, dilation_patchW, 
+          dH, dW,
+          n);
+    }
+  }));
+
+  return {gradInput1, gradInput2};
+}

aot/Pytorch-Correlation-extension/Correlation_Module/correlation_sampler.cpp

@@ -0,0 +1,138 @@
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <vector>
+#include <iostream>
+
+// declarations
+
+torch::Tensor correlation_cpp_forward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cpp_backward(
+    torch::Tensor grad_output,
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW);
+
+#ifdef USE_CUDA
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+#define CHECK_SAME_DEVICE(x, y) TORCH_CHECK(x.device() == y.device(), #x " is not on same device as " #y)
+
+torch::Tensor correlation_cuda_forward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cuda_backward(
+    torch::Tensor grad_output,
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW);
+
+// C++ interface
+
+torch::Tensor correlation_sample_forward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+  if (input1.device().is_cuda()){
+    CHECK_INPUT(input1);
+    CHECK_INPUT(input2);
+
+    // set device of input1 as default CUDA device
+    // https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_optional_c_u_d_a_guard.html
+    const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+    CHECK_SAME_DEVICE(input1, input2);
+
+    return correlation_cuda_forward(input1, input2, kH, kW, patchH, patchW,
+                             padH, padW, dilationH, dilationW,
+                             dilation_patchH, dilation_patchW,
+                             dH, dW);
+  }else{
+    return correlation_cpp_forward(input1, input2, kH, kW, patchH, patchW,
+                             padH, padW, dilationH, dilationW,
+                             dilation_patchH, dilation_patchW,
+                             dH, dW);
+  }
+}
+
+std::vector<torch::Tensor> correlation_sample_backward(
+    torch::Tensor input1,
+    torch::Tensor input2,
+    torch::Tensor grad_output,
+    int kH, int kW,
+    int patchH, int patchW,
+    int padH, int padW,
+    int dilationH, int dilationW,
+    int dilation_patchH, int dilation_patchW,
+    int dH, int dW) {
+
+  if(grad_output.device().is_cuda()){
+    CHECK_INPUT(input1);
+    CHECK_INPUT(input2);
+
+    // set device of input1 as default CUDA device
+    const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+    CHECK_SAME_DEVICE(input1, input2);
+    CHECK_SAME_DEVICE(input1, grad_output);
+
+    return correlation_cuda_backward(input1, input2, grad_output,
+                              kH, kW, patchH, patchW,
+                              padH, padW,
+                              dilationH, dilationW,
+                              dilation_patchH, dilation_patchW,
+                              dH, dW);
+  }else{
+    return correlation_cpp_backward(
+                              input1, input2, grad_output,
+                              kH, kW, patchH, patchW,
+                              padH, padW,
+                              dilationH, dilationW,
+                              dilation_patchH, dilation_patchW,
+                              dH, dW);
+  }
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &correlation_sample_forward, "Spatial Correlation Sampler Forward");
+  m.def("backward", &correlation_sample_backward, "Spatial Correlation Sampler backward");
+}
+
+#else
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &correlation_cpp_forward, "Spatial Correlation Sampler Forward");
+  m.def("backward", &correlation_cpp_backward, "Spatial Correlation Sampler backward");
+}
+
+#endif

aot/Pytorch-Correlation-extension/Correlation_Module/spatial_correlation_sampler/__init__.py

@@ -0,0 +1 @@
+from .spatial_correlation_sampler import SpatialCorrelationSampler, spatial_correlation_sample

aot/Pytorch-Correlation-extension/Correlation_Module/spatial_correlation_sampler/spatial_correlation_sampler.py

@@ -0,0 +1,107 @@
+from torch import nn
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.nn.modules.utils import _pair
+
+import spatial_correlation_sampler_backend as correlation
+
+
+def spatial_correlation_sample(input1,
+                               input2,
+                               kernel_size=1,
+                               patch_size=1,
+                               stride=1,
+                               padding=0,
+                               dilation=1,
+                               dilation_patch=1):
+    """Apply spatial correlation sampling on from input1 to input2,
+
+    Every parameter except input1 and input2 can be either single int
+    or a pair of int. For more information about Spatial Correlation
+    Sampling, see this page.
+    https://lmb.informatik.uni-freiburg.de/Publications/2015/DFIB15/
+
+    Args:
+        input1 : The first parameter.
+        input2 : The second parameter.
+        kernel_size : total size of your correlation kernel, in pixels
+        patch_size : total size of your patch, determining how many
+            different shifts will be applied
+        stride : stride of the spatial sampler, will modify output
+            height and width
+        padding : padding applied to input1 and input2 before applying
+            the correlation sampling, will modify output height and width
+        dilation_patch : step for every shift in patch
+
+    Returns:
+        Tensor: Result of correlation sampling
+
+    """
+    return SpatialCorrelationSamplerFunction.apply(input1, input2,
+                                                   kernel_size, patch_size,
+                                                   stride, padding, dilation, dilation_patch)
+
+
+class SpatialCorrelationSamplerFunction(Function):
+
+    @staticmethod
+    def forward(ctx,
+                input1,
+                input2,
+                kernel_size=1,
+                patch_size=1,
+                stride=1,
+                padding=0,
+                dilation=1,
+                dilation_patch=1):
+
+        ctx.save_for_backward(input1, input2)
+        kH, kW = ctx.kernel_size = _pair(kernel_size)
+        patchH, patchW = ctx.patch_size = _pair(patch_size)
+        padH, padW = ctx.padding = _pair(padding)
+        dilationH, dilationW = ctx.dilation = _pair(dilation)
+        dilation_patchH, dilation_patchW = ctx.dilation_patch = _pair(dilation_patch)
+        dH, dW = ctx.stride = _pair(stride)
+
+        output = correlation.forward(input1, input2,
+                                     kH, kW, patchH, patchW,
+                                     padH, padW, dilationH, dilationW,
+                                     dilation_patchH, dilation_patchW,
+                                     dH, dW)
+
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        input1, input2 = ctx.saved_variables
+
+        kH, kW = ctx.kernel_size
+        patchH, patchW = ctx.patch_size
+        padH, padW = ctx.padding
+        dilationH, dilationW = ctx.dilation
+        dilation_patchH, dilation_patchW = ctx.dilation_patch
+        dH, dW = ctx.stride
+
+        grad_input1, grad_input2 = correlation.backward(input1, input2, grad_output,
+                                                        kH, kW, patchH, patchW,
+                                                        padH, padW, dilationH, dilationW,
+                                                        dilation_patchH, dilation_patchW,
+                                                        dH, dW)
+        return grad_input1, grad_input2, None, None, None, None, None, None
+
+
+class SpatialCorrelationSampler(nn.Module):
+    def __init__(self, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1):
+        super(SpatialCorrelationSampler, self).__init__()
+        self.kernel_size = kernel_size
+        self.patch_size = patch_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.dilation_patch = dilation_patch
+
+    def forward(self, input1, input2):
+        return SpatialCorrelationSamplerFunction.apply(input1, input2, self.kernel_size,
+                                                       self.patch_size, self.stride,
+                                                       self.padding, self.dilation, self.dilation_patch)

aot/Pytorch-Correlation-extension/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [year] [fullname]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

aot/Pytorch-Correlation-extension/README.md

@@ -0,0 +1,155 @@
+
+[![PyPI](https://img.shields.io/pypi/v/spatial-correlation-sampler.svg)](https://pypi.org/project/spatial-correlation-sampler/)
+
+
+# Pytorch Correlation module
+
+this is a custom C++/Cuda implementation of Correlation module, used e.g. in [FlowNetC](https://arxiv.org/abs/1504.06852)
+
+This [tutorial](http://pytorch.org/tutorials/advanced/cpp_extension.html) was used as a basis for implementation, as well as
+[NVIDIA's cuda code](https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package)
+
+- Build and Install C++ and CUDA extensions by executing `python setup.py install`,
+- Benchmark C++ vs. CUDA by running `python benchmark.py {cpu, cuda}`,
+- Run gradient checks on the code by running `python grad_check.py --backend {cpu, cuda}`.
+
+# Requirements
+
+This module is expected to compile for Pytorch `2.1.0`.
+
+Before installation please check compatibility of your GPU and CUDA (_Compute Capability_) [nvidia docs](https://developer.nvidia.com/cuda-gpus). 
+e.g RTX 6000 is using CC=8.9 so we are setting the environment variable to
+
+`export TORCH_CUDA_ARCH_LIST="8.9+PTX"`
+
+# Installation
+
+be reminded this module requires `python3-dev` to compile C++ code, e.g. on Ubuntu run:
+
+`apt install python3-dev`
+
+this module is available on pip
+
+`pip install spatial-correlation-sampler`
+
+For a cpu-only version, you can install from source with
+
+`python setup_cpu.py install`
+
+# Known Problems
+
+This module needs compatible gcc version and CUDA to be compiled.
+Namely, CUDA 9.1 and below will need gcc5, while CUDA 9.2 and 10.0 will need gcc7
+See [this issue](https://github.com/ClementPinard/Pytorch-Correlation-extension/issues/1) for more information
+
+# Usage
+
+API has a few difference with NVIDIA's module
+ * output is now a 5D tensor, which reflects the shifts horizontal and vertical.
+ ```
+input (B x C x H x W) -> output (B x PatchH x PatchW x oH x oW)
+ ```
+ * Output sizes `oH` and `oW` are no longer dependant of patch size, but only of kernel size and padding
+ * Patch size `patch_size` is now the whole patch, and not only the radii.
+ * `stride1` is now `stride` and`stride2` is `dilation_patch`, which behave like dilated convolutions
+ * equivalent `max_displacement` is then `dilation_patch * (patch_size - 1) / 2`.
+ * `dilation` is a new parameter, it acts the same way as dilated convolution regarding the correlation kernel
+ * to get the right parameters for FlowNetC, you would have
+ ```
+kernel_size=1
+patch_size=21,
+stride=1,
+padding=0,
+dilation=1
+dilation_patch=2
+ ```
+
+
+## Example
+```python
+import torch
+from spatial_correlation_sampler import SpatialCorrelationSampler, 
+
+device = "cuda"
+batch_size = 1
+channel = 1
+H = 10
+W = 10
+dtype = torch.float32
+
+input1 = torch.randint(1, 4, (batch_size, channel, H, W), dtype=dtype, device=device, requires_grad=True)
+input2 = torch.randint_like(input1, 1, 4).requires_grad_(True)
+
+#You can either use the function or the module. Note that the module doesn't contain any parameter tensor.
+
+#function
+
+out = spatial_correlation_sample(input1,
+	                         input2,
+                                 kernel_size=3,
+                                 patch_size=1,
+                                 stride=2,
+                                 padding=0,
+                                 dilation=2,
+                                 dilation_patch=1)
+
+#module
+
+correlation_sampler = SpatialCorrelationSampler(
+    kernel_size=3,
+    patch_size=1,
+    stride=2,
+    padding=0,
+    dilation=2,
+    dilation_patch=1)
+out = correlation_sampler(input1, input2)
+
+```
+
+# Benchmark
+
+ * default parameters are from `benchmark.py`, FlowNetC parameters are same as use in `FlowNetC` with a batch size of 4, described in [this paper](https://arxiv.org/abs/1504.06852), implemented [here](https://github.com/lmb-freiburg/flownet2) and [here](https://github.com/NVIDIA/flownet2-pytorch/blob/master/networks/FlowNetC.py).
+ * Feel free to file an issue to add entries to this with your hardware !
+
+## CUDA Benchmark
+
+ * See [here](https://gist.github.com/ClementPinard/270e910147119831014932f67fb1b5ea) for a benchmark script working with [NVIDIA](https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package)'s code, and Pytorch.
+ * Benchmark are launched with environment variable `CUDA_LAUNCH_BLOCKING` set to `1`.
+ * Only `float32` is benchmarked.
+ * FlowNetC correlation parameters where launched with the following command:
+ 
+ ```bash
+ CUDA_LAUNCH_BLOCKING=1 python benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256 cuda -d float
+ 
+ CUDA_LAUNCH_BLOCKING=1 python NV_correlation_benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256
+ ```
+
+ | implementation | Correlation parameters |  device |     pass |      min time |      avg time |
+ | -------------- | ---------------------- | ------- | -------- | ------------: | ------------: |
+ |           ours |                default | 980 GTX |  forward |  **5.745 ms** |  **5.851 ms** |
+ |           ours |                default | 980 GTX | backward |     77.694 ms |     77.957 ms |
+ |         NVIDIA |                default | 980 GTX |  forward |     13.779 ms |     13.853 ms |
+ |         NVIDIA |                default | 980 GTX | backward | **73.383 ms** | **73.708 ms** |
+ |                |                        |         |          |               |               |
+ |           ours |               FlowNetC | 980 GTX |  forward |  **26.102 ms** |  **26.179 ms** |
+ |           ours |               FlowNetC | 980 GTX | backward | **208.091 ms** | **208.510 ms** |
+ |         NVIDIA |               FlowNetC | 980 GTX |  forward |      35.363 ms |      35.550 ms |
+ |         NVIDIA |               FlowNetC | 980 GTX | backward |     283.748 ms |     284.346 ms |
+ 
+### Notes
+ * The overhead of our implementation regarding `kernel_size` > 1 during backward needs some investigation, feel free to
+ dive in the code to improve it !
+ * The backward pass of NVIDIA is not entirely correct when stride1 > 1 and kernel_size > 1, because not everything
+ is computed, see [here](https://github.com/NVIDIA/flownet2-pytorch/blob/master/networks/correlation_package/src/correlation_cuda_kernel.cu#L120).
+
+## CPU Benchmark
+
+  * No other implementation is avalaible on CPU.
+  * It is obviously not recommended to run it on CPU if you have a GPU.
+
+ | Correlation parameters |               device |     pass |    min time |    avg time |
+ | ---------------------- | -------------------- | -------- | ----------: | ----------: |
+ |                default | E5-2630 v3 @ 2.40GHz |  forward |  159.616 ms |  188.727 ms |
+ |                default | E5-2630 v3 @ 2.40GHz | backward |  282.641 ms |  294.194 ms |
+ |               FlowNetC | E5-2630 v3 @ 2.40GHz |  forward |  2.138 s |  2.144 s |
+ |               FlowNetC | E5-2630 v3 @ 2.40GHz | backward | 7.006 s | 7.075 s |

aot/Pytorch-Correlation-extension/benchmark.py

@@ -0,0 +1,90 @@
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import time
+
+import torch
+from spatial_correlation_sampler import SpatialCorrelationSampler
+from tqdm import trange
+
+TIME_SCALES = {'s': 1, 'ms': 1000, 'us': 1000000}
+
+parser = argparse.ArgumentParser()
+parser.add_argument('backend', choices=['cpu', 'cuda'], default='cuda')
+parser.add_argument('-b', '--batch-size', type=int, default=16)
+parser.add_argument('-k', '--kernel-size', type=int, default=3)
+parser.add_argument('--patch', type=int, default=3)
+parser.add_argument('--patch_dilation', type=int, default=2)
+parser.add_argument('-c', '--channel', type=int, default=64)
+parser.add_argument('--height', type=int, default=100)
+parser.add_argument('-w', '--width', type=int, default=100)
+parser.add_argument('-s', '--stride', type=int, default=2)
+parser.add_argument('-p', '--pad', type=int, default=1)
+parser.add_argument('--scale', choices=['s', 'ms', 'us'], default='us')
+parser.add_argument('-r', '--runs', type=int, default=100)
+parser.add_argument('--dilation', type=int, default=2)
+parser.add_argument('-d', '--dtype', choices=['half', 'float', 'double'])
+
+args = parser.parse_args()
+
+device = torch.device(args.backend)
+
+if args.dtype == 'half':
+    dtype = torch.float16
+elif args.dtype == 'float':
+    dtype = torch.float32
+else:
+    dtype = torch.float64
+
+
+input1 = torch.randn(args.batch_size,
+                     args.channel,
+                     args.height,
+                     args.width,
+                     dtype=dtype,
+                     device=device,
+                     requires_grad=True)
+input2 = torch.randn_like(input1)
+
+correlation_sampler = SpatialCorrelationSampler(
+    args.kernel_size,
+    args.patch,
+    args.stride,
+    args.pad,
+    args.dilation,
+    args.patch_dilation)
+
+# Force CUDA initialization
+output = correlation_sampler(input1, input2)
+print(output.size())
+output.mean().backward()
+forward_min = float('inf')
+forward_time = 0
+backward_min = float('inf')
+backward_time = 0
+for _ in trange(args.runs):
+    correlation_sampler.zero_grad()
+
+    start = time.time()
+    output = correlation_sampler(input1, input2)
+    elapsed = time.time() - start
+    forward_min = min(forward_min, elapsed)
+    forward_time += elapsed
+    output = output.mean()
+
+    start = time.time()
+    (output.mean()).backward()
+    elapsed = time.time() - start
+    backward_min = min(backward_min, elapsed)
+    backward_time += elapsed
+
+scale = TIME_SCALES[args.scale]
+forward_min *= scale
+backward_min *= scale
+forward_average = forward_time / args.runs * scale
+backward_average = backward_time / args.runs * scale
+
+print('Forward: {0:.3f}/{1:.3f} {4} | Backward {2:.3f}/{3:.3f} {4}'.format(
+    forward_min, forward_average, backward_min, backward_average,
+    args.scale))

aot/Pytorch-Correlation-extension/check.py

@@ -0,0 +1,119 @@
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import numpy as np
+import torch
+
+from spatial_correlation_sampler import SpatialCorrelationSampler
+
+
+def check_equal(first, second, verbose):
+    if verbose:
+        print()
+    for i, (x, y) in enumerate(zip(first, second)):
+        x = x.cpu().detach().numpy()
+        y = y.cpu().detach().numpy()
+        if verbose:
+            print("x = {}".format(x.flatten()))
+            print("y = {}".format(y.flatten()))
+            print('-' * 80)
+        np.testing.assert_allclose(x, y, err_msg="Index: {}".format(i))
+
+
+def zero_grad(variables):
+    for variable in variables:
+        if variable.grad is not None: variable.grad.zero_()
+
+
+def get_grads(variables):
+    return [var.grad.clone() for var in variables]
+
+
+def check_forward(input1, input2, correlation_sampler, verbose, gpu_index=0):
+    device = torch.device(f"cuda:{gpu_index}")
+
+    cpu_values = correlation_sampler(input1, input2)
+    cuda_values = correlation_sampler(input1.to(device), input2.to(device))
+
+    print(f"Forward: CPU vs. CUDA device:{gpu_index} ... ", end='')
+    check_equal(cpu_values, cuda_values, verbose)
+    print('Ok')
+
+
+def check_backward(input1, input2, correlation_sampler, verbose, gpu_index=0):
+    device = torch.device(f"cuda:{gpu_index}")
+
+    zero_grad([input1, input2])
+
+    cpu_values = correlation_sampler(input1, input2)
+    cpu_values.sum().backward()
+    grad_cpu = get_grads([input1, input2])
+
+    zero_grad([input1, input2])
+
+    cuda_values = correlation_sampler(input1.to(device), input2.to(device))
+    cuda_values.sum().backward()
+    grad_cuda = get_grads([input1, input2])
+
+    print(f"Backward: CPU vs. CUDA device:{gpu_index} ... ", end='')
+    check_equal(grad_cpu, grad_cuda, verbose)
+    print('Ok')
+
+
+def check_multi_gpu_forward(correlation_sampler, verbose):
+    print("Multi-GPU forward")
+    total_gpus = torch.cuda.device_count()
+    for gpu in range(total_gpus):
+        check_forward(input1, input2, correlation_sampler, verbose, gpu_index=gpu)
+
+def check_multi_gpu_backward(correlation_sampler, verbose):
+    print("Multi-GPU backward")
+    total_gpus = torch.cuda.device_count()
+    for gpu in range(total_gpus):
+        check_backward(input1, input2, correlation_sampler, verbose, gpu_index=gpu)
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('direction', choices=['forward', 'backward'], nargs='+')
+parser.add_argument('-b', '--batch-size', type=int, default=1)
+parser.add_argument('-k', '--kernel-size', type=int, default=3)
+parser.add_argument('--patch', type=int, default=3)
+parser.add_argument('--patch_dilation', type=int, default=2)
+parser.add_argument('-c', '--channel', type=int, default=10)
+parser.add_argument('--height', type=int, default=10)
+parser.add_argument('-w', '--width', type=int, default=10)
+parser.add_argument('-s', '--stride', type=int, default=2)
+parser.add_argument('-p', '--pad', type=int, default=5)
+parser.add_argument('-v', '--verbose', action='store_true', default=False)
+parser.add_argument('-d', '--dilation', type=int, default=2)
+args = parser.parse_args()
+print(args)
+
+assert(torch.cuda.is_available()), "no comparison to make"
+input1 = torch.randn(args.batch_size,
+                     args.channel,
+                     args.height,
+                     args.width).double()
+input2 = torch.randn(args.batch_size,
+                     args.channel,
+                     args.height,
+                     args.width).double()
+input1.requires_grad = True
+input2.requires_grad = True
+
+correlation_sampler = SpatialCorrelationSampler(
+    args.kernel_size,
+    args.patch,
+    args.stride,
+    args.pad,
+    args.dilation,
+    args.patch_dilation)
+
+if 'forward' in args.direction:
+    check_forward(input1, input2, correlation_sampler, args.verbose)
+    if torch.cuda.device_count() > 1: check_multi_gpu_forward(correlation_sampler, args.verbose)
+
+if 'backward' in args.direction:
+    check_backward(input1, input2, correlation_sampler, args.verbose)
+    if torch.cuda.device_count() > 1: check_multi_gpu_backward(correlation_sampler, args.verbose)

aot/Pytorch-Correlation-extension/grad_check.py

@@ -0,0 +1,47 @@
+import argparse
+import torch
+# torch.set_printoptions(precision=1, threshold=10000)
+from torch.autograd import gradcheck
+from spatial_correlation_sampler import SpatialCorrelationSampler
+
+parser = argparse.ArgumentParser()
+parser.add_argument('backend', choices=['cpu', 'cuda'], default='cuda')
+parser.add_argument('-b', '--batch-size', type=int, default=2)
+parser.add_argument('-k', '--kernel-size', type=int, default=3)
+parser.add_argument('--patch', type=int, default=3)
+parser.add_argument('--patch_dilation', type=int, default=2)
+parser.add_argument('-c', '--channel', type=int, default=2)
+parser.add_argument('--height', type=int, default=10)
+parser.add_argument('-w', '--width', type=int, default=10)
+parser.add_argument('-s', '--stride', type=int, default=2)
+parser.add_argument('-p', '--pad', type=int, default=1)
+parser.add_argument('-d', '--dilation', type=int, default=2)
+
+args = parser.parse_args()
+
+input1 = torch.randn(args.batch_size,
+                     args.channel,
+                     args.height,
+                     args.width,
+                     dtype=torch.float64,
+                     device=torch.device(args.backend))
+input2 = torch.randn(args.batch_size,
+                     args.channel,
+                     args.height,
+                     args.width,
+                     dtype=torch.float64,
+                     device=torch.device(args.backend))
+
+input1.requires_grad = True
+input2.requires_grad = True
+
+correlation_sampler = SpatialCorrelationSampler(args.kernel_size,
+                                                args.patch,
+                                                args.stride,
+                                                args.pad,
+                                                args.dilation,
+                                                args.patch_dilation)
+
+
+if gradcheck(correlation_sampler, [input1, input2]):
+    print('Ok')

aot/Pytorch-Correlation-extension/requirements.txt

@@ -0,0 +1,2 @@
+torch>=1.0.1
+numpy

aot/Pytorch-Correlation-extension/setup.py

@@ -0,0 +1,69 @@
+import os
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
+from os.path import join
+
+CPU_ONLY = False
+project_root = 'Correlation_Module'
+
+source_files = ['correlation.cpp', 'correlation_sampler.cpp']
+
+cxx_args = ['-std=c++17', '-fopenmp']
+
+def generate_nvcc_args(gpu_archs):
+    nvcc_args = []
+    for arch in gpu_archs:
+        nvcc_args.extend(['-gencode', f'arch=compute_{arch},code=sm_{arch}'])
+    return nvcc_args
+
+gpu_arch = os.environ.get('GPU_ARCH', '').split()
+nvcc_args = generate_nvcc_args(gpu_arch)
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+
+def launch_setup():
+    if CPU_ONLY:
+        Extension = CppExtension
+        macro = []
+    else:
+        Extension = CUDAExtension
+        source_files.append('correlation_cuda_kernel.cu')
+        macro = [("USE_CUDA", None)]
+
+    sources = [join(project_root, file) for file in source_files]
+
+    setup(
+        name='spatial_correlation_sampler',
+        version="0.4.0",
+        author="Clément Pinard",
+        author_email="[email protected]",
+        description="Correlation module for pytorch",
+        long_description=long_description,
+        long_description_content_type="text/markdown",
+        url="https://github.com/ClementPinard/Pytorch-Correlation-extension",
+        install_requires=['torch>=1.1', 'numpy'],
+        ext_modules=[
+            Extension('spatial_correlation_sampler_backend',
+                      sources,
+                      define_macros=macro,
+                      extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args},
+                      extra_link_args=['-lgomp'])
+        ],
+        package_dir={'': project_root},
+        packages=['spatial_correlation_sampler'],
+        cmdclass={
+            'build_ext': BuildExtension
+        },
+        classifiers=[
+            "Programming Language :: Python :: 3",
+            "License :: OSI Approved :: MIT License",
+            "Operating System :: POSIX :: Linux",
+            "Intended Audience :: Science/Research",
+            "Topic :: Scientific/Engineering :: Artificial Intelligence"
+        ])
+
+
+if __name__ == '__main__':
+    launch_setup()

aot/Pytorch-Correlation-extension/setup_cpu.py

@@ -0,0 +1,4 @@
+import setup
+
+setup.CPU_ONLY = True
+setup.launch_setup()

aot/README.md

@@ -0,0 +1,152 @@
+# AOT Series Frameworks in PyTorch
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/decoupling-features-in-hierarchical/semi-supervised-video-object-segmentation-on-15)](https://paperswithcode.com/sota/semi-supervised-video-object-segmentation-on-15?p=decoupling-features-in-hierarchical)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/associating-objects-with-scalable/video-object-segmentation-on-youtube-vos)](https://paperswithcode.com/sota/video-object-segmentation-on-youtube-vos?p=associating-objects-with-scalable)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/associating-objects-with-scalable/semi-supervised-video-object-segmentation-on-18)](https://paperswithcode.com/sota/semi-supervised-video-object-segmentation-on-18?p=associating-objects-with-scalable)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/associating-objects-with-scalable/semi-supervised-video-object-segmentation-on-1)](https://paperswithcode.com/sota/semi-supervised-video-object-segmentation-on-1?p=associating-objects-with-scalable)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/associating-objects-with-scalable/visual-object-tracking-on-davis-2017)](https://paperswithcode.com/sota/visual-object-tracking-on-davis-2017?p=associating-objects-with-scalable)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/associating-objects-with-scalable/visual-object-tracking-on-davis-2016)](https://paperswithcode.com/sota/visual-object-tracking-on-davis-2016?p=associating-objects-with-scalable)
+
+A modular reference PyTorch implementation of AOT series frameworks:
+- **DeAOT**: Decoupling Features in Hierachical Propagation for Video Object Segmentation (NeurIPS 2022, Spotlight) [[OpenReview](https://openreview.net/forum?id=DgM7-7eMkq0)][[PDF](https://arxiv.org/pdf/2210.09782.pdf)]
+<img src="source/overview_deaot.png" width="90%"/>
+
+- **AOT**: Associating Objects with Transformers for Video Object Segmentation (NeurIPS 2021, Score 8/8/7/8) [[OpenReview](https://openreview.net/forum?id=hl3v8io3ZYt)][[PDF](https://arxiv.org/abs/2106.02638)]
+<img src="source/overview.png" width="90%"/>
+
+An extension of AOT, [AOST](https://arxiv.org/abs/2203.11442) (under review), is available now. AOST is a more robust and flexible framework, supporting run-time speed-accuracy trade-offs.
+
+## Examples
+Benchmark examples:
+
+<img src="source/some_results.png" width="81%"/>
+
+General examples (Messi and Kobe):
+
+<img src="source/messi.gif" width="45%"/> <img src="source/kobe.gif" width="45%"/>
+
+## Highlights
+- **High performance:** up to **85.5%** ([R50-AOTL](MODEL_ZOO.md#youtube-vos-2018-val)) on YouTube-VOS 2018 and **82.1%** ([SwinB-AOTL]((MODEL_ZOO.md#youtube-vos-2018-val))) on DAVIS-2017 Test-dev under standard settings (without any test-time augmentation and post processing). 
+- **High efficiency:** up to **51fps** ([AOTT](MODEL_ZOO.md#davis-2017-test)) on DAVIS-2017 (480p) even with **10** objects and **41fps** on YouTube-VOS (1.3x480p). AOT can process multiple objects (less than a pre-defined number, 10 is the default) as efficiently as processing a single object. This project also supports inferring any number of objects together within a video by automatic separation and aggregation.
+- **Multi-GPU training and inference**
+- **Mixed precision training and inference** 
+- **Test-time augmentation:** multi-scale and flipping augmentations are supported.
+
+## Requirements
+   * Python3
+   * pytorch >= 1.7.0 and torchvision
+   * opencv-python
+   * Pillow
+   * Pytorch Correlation (Recommend to install from [source](https://github.com/ClementPinard/Pytorch-Correlation-extension) instead of using `pip`. **The project can also work without this module but will lose some efficiency of the short-term attention**.)
+
+Optional:
+   * scikit-image (if you want to run our **Demo**, please install)
+
+## Model Zoo and Results
+Pre-trained models, benckmark scores, and pre-computed results reproduced by this project can be found in [MODEL_ZOO.md](MODEL_ZOO.md).
+
+## Demo - Panoptic Propagation
+We provide a simple demo to demonstrate AOT's effectiveness. The demo will propagate more than **40** objects, including semantic regions (like sky) and instances (like person), together within a single complex scenario and predict its video panoptic segmentation.
+
+To run the demo, download the [checkpoint](https://drive.google.com/file/d/1qJDYn3Ibpquu4ffYoQmVjg1YCbr2JQep/view?usp=sharing) of R50-AOTL into [pretrain_models](pretrain_models), and then run:
+```bash
+python tools/demo.py
+```
+which will predict the given scenarios in the resolution of 1.3x480p. You can also run this demo with other AOTs ([MODEL_ZOO.md](MODEL_ZOO.md)) by setting `--model` (model type) and `--ckpt_path` (checkpoint path).
+
+Two scenarios from [VSPW](https://www.vspwdataset.com/home) are supplied in [datasets/Demo](datasets/Demo):
+
+- 1001_3iEIq5HBY1s: 44 objects. 1080P.
+- 1007_YCTBBdbKSSg: 43 objects. 1080P.
+
+Results:
+
+<img src="source/1001_3iEIq5HBY1s.gif" width="45%"/>  <img src="source/1007_YCTBBdbKSSg.gif" width="45%"/>
+
+
+## Getting Started
+0. Prepare a valid environment follow the [requirements](#requirements).
+
+1. Prepare datasets:
+
+    Please follow the below instruction to prepare datasets in each corresponding folder.
+    * **Static** 
+    
+        [datasets/Static](datasets/Static): pre-training dataset with static images. Guidance can be found in [AFB-URR](https://github.com/xmlyqing00/AFB-URR), which we referred to in the implementation of the pre-training.
+    * **YouTube-VOS**
+
+        A commonly-used large-scale VOS dataset.
+
+        [datasets/YTB/2019](datasets/YTB/2019): version 2019, download [link](https://drive.google.com/drive/folders/1BWzrCWyPEmBEKm0lOHe5KLuBuQxUSwqz?usp=sharing). `train` is required for training. `valid` (6fps) and `valid_all_frames` (30fps, optional) are used for evaluation.
+
+        [datasets/YTB/2018](datasets/YTB/2018): version 2018, download [link](https://drive.google.com/drive/folders/1bI5J1H3mxsIGo7Kp-pPZU8i6rnykOw7f?usp=sharing). Only `valid` (6fps) and `valid_all_frames` (30fps, optional) are required for this project and used for evaluation.
+
+    * **DAVIS**
+
+        A commonly-used small-scale VOS dataset.
+
+        [datasets/DAVIS](datasets/DAVIS): [TrainVal](https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-trainval-480p.zip) (480p) contains both the training and validation split. [Test-Dev](https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-test-dev-480p.zip) (480p) contains the Test-dev split. The [full-resolution version](https://davischallenge.org/davis2017/code.html) is also supported for training and evaluation but not required.
+
+
+2. Prepare ImageNet pre-trained encoders
+
+    Select and download below checkpoints into [pretrain_models](pretrain_models):
+
+    - [MobileNet-V2](https://download.pytorch.org/models/mobilenet_v2-b0353104.pth) (default encoder)
+    - [MobileNet-V3](https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth)
+    - [ResNet-50](https://download.pytorch.org/models/resnet50-0676ba61.pth)
+    - [ResNet-101](https://download.pytorch.org/models/resnet101-63fe2227.pth)
+    - [ResNeSt-50](https://github.com/zhanghang1989/ResNeSt/releases/download/weights_step1/resnest50-528c19ca.pth)
+    - [ResNeSt-101](https://github.com/zhanghang1989/ResNeSt/releases/download/weights_step1/resnest101-22405ba7.pth)
+    - [Swin-Base](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth)
+
+    The current default training configs are not optimized for encoders larger than ResNet-50. If you want to use larger encoders, we recommend early stopping the main-training stage at 80,000 iterations (100,000 in default) to avoid over-fitting on the seen classes of YouTube-VOS.
+
+
+
+3. Training and Evaluation
+
+    The [example script](train_eval.sh) will train AOTT with 2 stages using 4 GPUs and auto-mixed precision (`--amp`). The first stage is a pre-training stage using `Static` dataset, and the second stage is a main-training stage, which uses both `YouTube-VOS 2019 train` and `DAVIS-2017 train` for training, resulting in a model that can generalize to different domains (YouTube-VOS and DAVIS) and different frame rates (6fps, 24fps, and 30fps).
+
+    Notably, you can use only the `YouTube-VOS 2019 train` split in the second stage by changing `pre_ytb_dav` to `pre_ytb`, which leads to better YouTube-VOS performance on unseen classes. Besides, if you don't want to do the first stage, you can start the training from stage `ytb`, but the performance will drop about 1~2% absolutely.
+
+    After the training is finished (about 0.6 days for each stage with 4 Tesla V100 GPUs), the [example script](train_eval.sh) will evaluate the model on YouTube-VOS and DAVIS, and the results will be packed into Zip files. For calculating scores, please use official YouTube-VOS servers ([2018 server](https://competitions.codalab.org/competitions/19544) and [2019 server](https://competitions.codalab.org/competitions/20127)), official [DAVIS toolkit](https://github.com/davisvideochallenge/davis-2017) (for Val), and official [DAVIS server](https://competitions.codalab.org/competitions/20516#learn_the_details) (for Test-dev).
+
+## Adding your own dataset
+Coming
+
+## Troubleshooting
+Waiting
+
+## TODO
+- [ ] Code documentation
+- [ ] Adding your own dataset
+- [ ] Results with test-time augmentations in Model Zoo
+- [ ] Support gradient accumulation
+- [x] Demo tool
+
+## Citations
+Please consider citing the related paper(s) in your publications if it helps your research.
+```
+@inproceedings{yang2022deaot,
+  title={Decoupling Features in Hierarchical Propagation for Video Object Segmentation},
+  author={Yang, Zongxin and Yang, Yi},
+  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
+  year={2022}
+}
+@article{yang2021aost,
+  title={Scalable Multi-object Identification for Video Object Segmentation},
+  author={Yang, Zongxin and Miao, Jiaxu and Wang, Xiaohan and Wei, Yunchao and Yang, Yi},
+  journal={arXiv preprint arXiv:2203.11442},
+  year={2022}
+}
+@inproceedings{yang2021aot,
+  title={Associating Objects with Transformers for Video Object Segmentation},
+  author={Yang, Zongxin and Wei, Yunchao and Yang, Yi},
+  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
+  year={2021}
+}
+```
+
+## License
+This project is released under the BSD-3-Clause license. See [LICENSE](LICENSE) for additional details.

aot/configs/default.py

@@ -0,0 +1,138 @@
+import os
+import importlib
+
+
+class DefaultEngineConfig():
+    def __init__(self, exp_name='default', model='aott'):
+        model_cfg = importlib.import_module('configs.models.' +
+                                            model).ModelConfig()
+        self.__dict__.update(model_cfg.__dict__)  # add model config
+
+        self.EXP_NAME = exp_name + '_' + self.MODEL_NAME
+
+        self.STAGE_NAME = 'YTB'
+
+        self.DATASETS = ['youtubevos']
+        self.DATA_WORKERS = 8
+        self.DATA_RANDOMCROP = (465,
+                                465) if self.MODEL_ALIGN_CORNERS else (464,
+                                                                       464)
+        self.DATA_RANDOMFLIP = 0.5
+        self.DATA_MAX_CROP_STEPS = 10
+        self.DATA_SHORT_EDGE_LEN = 480
+        self.DATA_MIN_SCALE_FACTOR = 0.7
+        self.DATA_MAX_SCALE_FACTOR = 1.3
+        self.DATA_RANDOM_REVERSE_SEQ = True
+        self.DATA_SEQ_LEN = 5
+        self.DATA_DAVIS_REPEAT = 5
+        self.DATA_RANDOM_GAP_DAVIS = 12  # max frame interval between two sampled frames for DAVIS (24fps)
+        self.DATA_RANDOM_GAP_YTB = 3  # max frame interval between two sampled frames for YouTube-VOS (6fps)
+        self.DATA_DYNAMIC_MERGE_PROB = 0.3
+
+        self.PRETRAIN = True
+        self.PRETRAIN_FULL = False  # if False, load encoder only
+        self.PRETRAIN_MODEL = './data_wd/pretrain_model/mobilenet_v2.pth'
+        # self.PRETRAIN_MODEL = './pretrain_models/mobilenet_v2-b0353104.pth'
+
+        self.TRAIN_TOTAL_STEPS = 100000
+        self.TRAIN_START_STEP = 0
+        self.TRAIN_WEIGHT_DECAY = 0.07
+        self.TRAIN_WEIGHT_DECAY_EXCLUSIVE = {
+            # 'encoder.': 0.01
+        }
+        self.TRAIN_WEIGHT_DECAY_EXEMPTION = [
+            'absolute_pos_embed', 'relative_position_bias_table',
+            'relative_emb_v', 'conv_out'
+        ]
+        self.TRAIN_LR = 2e-4
+        self.TRAIN_LR_MIN = 2e-5 if 'mobilenetv2' in self.MODEL_ENCODER else 1e-5
+        self.TRAIN_LR_POWER = 0.9
+        self.TRAIN_LR_ENCODER_RATIO = 0.1
+        self.TRAIN_LR_WARM_UP_RATIO = 0.05
+        self.TRAIN_LR_COSINE_DECAY = False
+        self.TRAIN_LR_RESTART = 1
+        self.TRAIN_LR_UPDATE_STEP = 1
+        self.TRAIN_AUX_LOSS_WEIGHT = 1.0
+        self.TRAIN_AUX_LOSS_RATIO = 1.0
+        self.TRAIN_OPT = 'adamw'
+        self.TRAIN_SGD_MOMENTUM = 0.9
+        self.TRAIN_GPUS = 4
+        self.TRAIN_BATCH_SIZE = 16
+        self.TRAIN_TBLOG = False
+        self.TRAIN_TBLOG_STEP = 50
+        self.TRAIN_LOG_STEP = 20
+        self.TRAIN_IMG_LOG = True
+        self.TRAIN_TOP_K_PERCENT_PIXELS = 0.15
+        self.TRAIN_SEQ_TRAINING_FREEZE_PARAMS = ['patch_wise_id_bank']
+        self.TRAIN_SEQ_TRAINING_START_RATIO = 0.5
+        self.TRAIN_HARD_MINING_RATIO = 0.5
+        self.TRAIN_EMA_RATIO = 0.1
+        self.TRAIN_CLIP_GRAD_NORM = 5.
+        self.TRAIN_SAVE_STEP = 5000
+        self.TRAIN_MAX_KEEP_CKPT = 8
+        self.TRAIN_RESUME = False
+        self.TRAIN_RESUME_CKPT = None
+        self.TRAIN_RESUME_STEP = 0
+        self.TRAIN_AUTO_RESUME = True
+        self.TRAIN_DATASET_FULL_RESOLUTION = False
+        self.TRAIN_ENABLE_PREV_FRAME = False
+        self.TRAIN_ENCODER_FREEZE_AT = 2
+        self.TRAIN_LSTT_EMB_DROPOUT = 0.
+        self.TRAIN_LSTT_ID_DROPOUT = 0.
+        self.TRAIN_LSTT_DROPPATH = 0.1
+        self.TRAIN_LSTT_DROPPATH_SCALING = False
+        self.TRAIN_LSTT_DROPPATH_LST = False
+        self.TRAIN_LSTT_LT_DROPOUT = 0.
+        self.TRAIN_LSTT_ST_DROPOUT = 0.
+
+        self.TEST_GPU_ID = 0
+        self.TEST_GPU_NUM = 1
+        self.TEST_FRAME_LOG = False
+        self.TEST_DATASET = 'youtubevos'
+        self.TEST_DATASET_FULL_RESOLUTION = False
+        self.TEST_DATASET_SPLIT = 'val'
+        self.TEST_CKPT_PATH = None
+        # if "None", evaluate the latest checkpoint.
+        self.TEST_CKPT_STEP = None
+        self.TEST_FLIP = False
+        self.TEST_MULTISCALE = [1]
+        self.TEST_MAX_SHORT_EDGE = None
+        self.TEST_MAX_LONG_EDGE = 800 * 1.3
+        self.TEST_WORKERS = 4
+
+        # GPU distribution
+        self.DIST_ENABLE = True
+        self.DIST_BACKEND = "nccl"  # "gloo"
+        self.DIST_URL = "tcp://127.0.0.1:13241"
+        self.DIST_START_GPU = 0
+
+    def init_dir(self):
+        self.DIR_DATA = '../VOS02/datasets'#'./datasets'
+        self.DIR_DAVIS = os.path.join(self.DIR_DATA, 'DAVIS')
+        self.DIR_YTB = os.path.join(self.DIR_DATA, 'YTB')
+        self.DIR_STATIC = os.path.join(self.DIR_DATA, 'Static')
+
+        self.DIR_ROOT = './'#'./data_wd/youtube_vos_jobs'
+
+        self.DIR_RESULT = os.path.join(self.DIR_ROOT, 'result', self.EXP_NAME,
+                                       self.STAGE_NAME)
+        self.DIR_CKPT = os.path.join(self.DIR_RESULT, 'ckpt')
+        self.DIR_EMA_CKPT = os.path.join(self.DIR_RESULT, 'ema_ckpt')
+        self.DIR_LOG = os.path.join(self.DIR_RESULT, 'log')
+        self.DIR_TB_LOG = os.path.join(self.DIR_RESULT, 'log', 'tensorboard')
+        # self.DIR_IMG_LOG = os.path.join(self.DIR_RESULT, 'log', 'img')
+        # self.DIR_EVALUATION = os.path.join(self.DIR_RESULT, 'eval')
+        self.DIR_IMG_LOG = './img_logs'
+        self.DIR_EVALUATION = './results'
+
+        for path in [
+                self.DIR_RESULT, self.DIR_CKPT, self.DIR_EMA_CKPT,
+                self.DIR_LOG, self.DIR_EVALUATION, self.DIR_IMG_LOG,
+                self.DIR_TB_LOG
+        ]:
+            if not os.path.isdir(path):
+                try:
+                    os.makedirs(path)
+                except Exception as inst:
+                    print(inst)
+                    print('Failed to make dir: {}.'.format(path))

aot/configs/models/aotb.py

@@ -0,0 +1,9 @@
+import os
+from .default import DefaultModelConfig
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'AOTB'
+
+        self.MODEL_LSTT_NUM = 3

aot/configs/models/aotl.py

@@ -0,0 +1,13 @@
+import os
+from .default import DefaultModelConfig
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'AOTL'
+
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/aots.py

@@ -0,0 +1,9 @@
+import os
+from .default import DefaultModelConfig
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'AOTS'
+
+        self.MODEL_LSTT_NUM = 2

aot/configs/models/aott.py

@@ -0,0 +1,7 @@
+import os
+from .default import DefaultModelConfig
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'AOTT'

aot/configs/models/deaotb.py

@@ -0,0 +1,9 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'DeAOTB'
+
+        self.MODEL_LSTT_NUM = 3

aot/configs/models/deaotl.py

@@ -0,0 +1,13 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'DeAOTL'
+
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/deaots.py

@@ -0,0 +1,9 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'DeAOTS'
+
+        self.MODEL_LSTT_NUM = 2

aot/configs/models/deaott.py

@@ -0,0 +1,7 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'DeAOTT'

aot/configs/models/default.py

@@ -0,0 +1,27 @@
+class DefaultModelConfig():
+    def __init__(self):
+        self.MODEL_NAME = 'AOTDefault'
+
+        self.MODEL_VOS = 'aot'
+        self.MODEL_ENGINE = 'aotengine'
+        self.MODEL_ALIGN_CORNERS = True
+        self.MODEL_ENCODER = 'mobilenetv2'
+        self.MODEL_ENCODER_PRETRAIN = './pretrain_models/mobilenet_v2-b0353104.pth'
+        self.MODEL_ENCODER_DIM = [24, 32, 96, 1280]  # 4x, 8x, 16x, 16x
+        self.MODEL_ENCODER_EMBEDDING_DIM = 256
+        self.MODEL_DECODER_INTERMEDIATE_LSTT = True
+        self.MODEL_FREEZE_BN = True
+        self.MODEL_FREEZE_BACKBONE = False
+        self.MODEL_MAX_OBJ_NUM = 10
+        self.MODEL_SELF_HEADS = 8
+        self.MODEL_ATT_HEADS = 8
+        self.MODEL_LSTT_NUM = 1
+        self.MODEL_EPSILON = 1e-5
+        self.MODEL_USE_PREV_PROB = False
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 9999
+        self.TRAIN_AUG_TYPE = 'v1'
+
+        self.TEST_LONG_TERM_MEM_GAP = 9999
+
+        self.TEST_SHORT_TERM_MEM_SKIP = 1

aot/configs/models/default_deaot.py

@@ -0,0 +1,17 @@
+from .default import DefaultModelConfig as BaseConfig
+
+
+class DefaultModelConfig(BaseConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'DeAOTDefault'
+
+        self.MODEL_VOS = 'deaot'
+        self.MODEL_ENGINE = 'deaotengine'
+
+        self.MODEL_DECODER_INTERMEDIATE_LSTT = False
+
+        self.MODEL_SELF_HEADS = 1
+        self.MODEL_ATT_HEADS = 1
+
+        self.TRAIN_AUG_TYPE = 'v2'

aot/configs/models/r101_aotl.py

@@ -0,0 +1,16 @@
+from .default import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'R101_AOTL'
+
+        self.MODEL_ENCODER = 'resnet101'
+        self.MODEL_ENCODER_PRETRAIN = './pretrain_models/resnet101-63fe2227.pth'  # https://download.pytorch.org/models/resnet101-63fe2227.pth
+        self.MODEL_ENCODER_DIM = [256, 512, 1024, 1024]  # 4x, 8x, 16x, 16x
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/r50_aotl.py

@@ -0,0 +1,16 @@
+from .default import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'R50_AOTL'
+
+        self.MODEL_ENCODER = 'resnet50'
+        self.MODEL_ENCODER_PRETRAIN = './pretrain_models/resnet50-0676ba61.pth'  # https://download.pytorch.org/models/resnet50-0676ba61.pth
+        self.MODEL_ENCODER_DIM = [256, 512, 1024, 1024]  # 4x, 8x, 16x, 16x
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/r50_deaotl.py

@@ -0,0 +1,16 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'R50_DeAOTL'
+
+        self.MODEL_ENCODER = 'resnet50'
+        self.MODEL_ENCODER_DIM = [256, 512, 1024, 1024]  # 4x, 8x, 16x, 16x
+
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/rs101_aotl.py

@@ -0,0 +1,16 @@
+from .default import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'R101_AOTL'
+
+        self.MODEL_ENCODER = 'resnest101'
+        self.MODEL_ENCODER_PRETRAIN = './pretrain_models/resnest101-22405ba7.pth'  # https://github.com/zhanghang1989/ResNeSt/releases/download/weights_step1/resnest101-22405ba7.pth
+        self.MODEL_ENCODER_DIM = [256, 512, 1024, 1024]  # 4x, 8x, 16x, 16x
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/swinb_aotl.py

@@ -0,0 +1,17 @@
+from .default import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'SwinB_AOTL'
+
+        self.MODEL_ENCODER = 'swin_base'
+        self.MODEL_ENCODER_PRETRAIN = './pretrain_models/swin_base_patch4_window7_224_22k.pth'  # https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth
+        self.MODEL_ALIGN_CORNERS = False
+        self.MODEL_ENCODER_DIM = [128, 256, 512, 512]  # 4x, 8x, 16x, 16x
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/models/swinb_deaotl.py

@@ -0,0 +1,17 @@
+from .default_deaot import DefaultModelConfig
+
+
+class ModelConfig(DefaultModelConfig):
+    def __init__(self):
+        super().__init__()
+        self.MODEL_NAME = 'SwinB_DeAOTL'
+
+        self.MODEL_ENCODER = 'swin_base'
+        self.MODEL_ALIGN_CORNERS = False
+        self.MODEL_ENCODER_DIM = [128, 256, 512, 512]  # 4x, 8x, 16x, 16x
+
+        self.MODEL_LSTT_NUM = 3
+
+        self.TRAIN_LONG_TERM_MEM_GAP = 2
+
+        self.TEST_LONG_TERM_MEM_GAP = 5

aot/configs/pre.py

@@ -0,0 +1,19 @@
+from .default import DefaultEngineConfig
+
+
+class EngineConfig(DefaultEngineConfig):
+    def __init__(self, exp_name='default', model='AOTT'):
+        super().__init__(exp_name, model)
+        self.STAGE_NAME = 'PRE'
+
+        self.init_dir()
+
+        self.DATASETS = ['static']
+
+        self.DATA_DYNAMIC_MERGE_PROB = 1.0
+
+        self.TRAIN_LR = 4e-4
+        self.TRAIN_LR_MIN = 2e-5
+        self.TRAIN_WEIGHT_DECAY = 0.03
+        self.TRAIN_SEQ_TRAINING_START_RATIO = 1.0
+        self.TRAIN_AUX_LOSS_RATIO = 0.1

aot/configs/pre_dav.py

@@ -0,0 +1,21 @@
+import os
+from .default import DefaultEngineConfig
+
+
+class EngineConfig(DefaultEngineConfig):
+    def __init__(self, exp_name='default', model='AOTT'):
+        super().__init__(exp_name, model)
+        self.STAGE_NAME = 'PRE_DAV'
+
+        self.init_dir()
+
+        self.DATASETS = ['davis2017']
+
+        self.TRAIN_TOTAL_STEPS = 50000
+
+        pretrain_stage = 'PRE'
+        pretrain_ckpt = 'save_step_100000.pth'
+        self.PRETRAIN_FULL = True  # if False, load encoder only
+        self.PRETRAIN_MODEL = os.path.join(self.DIR_ROOT, 'result',
+                                           self.EXP_NAME, pretrain_stage,
+                                           'ema_ckpt', pretrain_ckpt)

aot/configs/pre_ytb.py

@@ -0,0 +1,17 @@
+import os
+from .default import DefaultEngineConfig
+
+
+class EngineConfig(DefaultEngineConfig):
+    def __init__(self, exp_name='default', model='AOTT'):
+        super().__init__(exp_name, model)
+        self.STAGE_NAME = 'PRE_YTB'
+
+        self.init_dir()
+
+        pretrain_stage = 'PRE'
+        pretrain_ckpt = 'save_step_100000.pth'
+        self.PRETRAIN_FULL = True  # if False, load encoder only
+        self.PRETRAIN_MODEL = os.path.join(self.DIR_ROOT, 'result',
+                                           self.EXP_NAME, pretrain_stage,
+                                           'ema_ckpt', pretrain_ckpt)