Upload ATCTrack-VLM code and selected checkpoints

.gitattributes

@@ -33,3 +33,8 @@ 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
+asset/experiment.png filter=lfs diff=lfs merge=lfs -text
+asset/framework.png filter=lfs diff=lfs merge=lfs -text
+asset/motivation.png filter=lfs diff=lfs merge=lfs -text
+resource/text_infor/refcoco_4o_v3_extract_mask.json filter=lfs diff=lfs merge=lfs -text
+resource/text_infor/vasttrack_4o_v3_extract_mask.json filter=lfs diff=lfs merge=lfs -text

QWEN_STATE_CURRENT_LOGIC.md

@@ -0,0 +1,264 @@
+# ATCTrack-VLM Qwen 更新器当前逻辑
+
+本文档按当前代码真实逻辑梳理，重点包括 Qwen 更新器、训练参数、loss、保存方式和具体代码位置。
+
+## 1. Qwen 更新器结构
+
+Qwen 更新器类位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:70`
+
+```python
+QwenTargetStateEncoder
+```
+
+初始化逻辑：
+
+- 加载 Qwen processor/tokenizer/model：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:98`
+- `padding_side` 设置为 `left`：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:104`
+- 新增 5 个 token：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:107`
+
+```text
+<answer>
+</answer>
+<state_token>
+</state_token>
+<TARGET_STATE>
+```
+
+- Qwen 主体先冻结：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:125`
+- 如果 `USE_LORA: true`，加 LoRA：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:132`
+- `<TARGET_STATE>` embedding 训练方式：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:151`
+
+现在不是外挂一个独立参数，而是训练 Qwen 原始 embedding 表里的 `<TARGET_STATE>` 那一行。通过 grad hook 只保留这一行梯度，其它 token embedding 梯度置 0：
+
+```python
+mask.index_fill_(0, token_ids, 1)
+return grad * mask.view(-1, 1)
+```
+
+## 2. Qwen 输入和生成流程
+
+prompt 位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:210`
+
+当前 prompt 里不会直接出现 `<TARGET_STATE>` 字符串，只说生成 special target-state token。
+
+forward 流程位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:439`
+
+```text
+template[-2] 图片 + bbox
+template[-1] 图片 + bbox
+prompt
+-> processor
+-> Qwen generate 一次
+-> 输出 <answer>yes/no</answer><state_token><TARGET_STATE></state_token>
+```
+
+具体流程：
+
+- 两张图片画 bbox：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:445`
+- 构造 image + text message：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:448`
+- processor 编码图文输入：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:462`
+- Qwen generate：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:464`
+
+生成被限制成两种合法格式之一：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:324`
+
+```text
+<answer>yes</answer><state_token><TARGET_STATE></state_token>
+<answer>no</answer><state_token><TARGET_STATE></state_token>
+```
+
+限制函数位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:348`
+
+核心是 `prefix_allowed_tokens_fn`。
+
+## 3. Token Hidden State 使用方式
+
+`<TARGET_STATE>` hidden state 提取位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:418`
+
+逻辑：
+
+```text
+如果生成了 <TARGET_STATE>：
+    取 <TARGET_STATE> 对应 hidden state
+否则：
+    fallback 到最后一个非 pad 生成 token 的 hidden state
+```
+
+然后：
+
+```text
+h_target -> projector -> z_target
+```
+
+代码位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:479`
+
+```python
+h_target = self._target_hidden_from_generation(...)
+z_target = self.projector(h_target)
+```
+
+`z_target` 注入 tracker 的方式是 FiLM 调制：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:485`
+
+```python
+opt_feat * (1 + gate * gamma) + gate * beta
+```
+
+实际注入位置在 tracker temporal loop：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:823`
+
+## 4. Template 更新逻辑
+
+训练时用 3 张 template、4 张 search：
+
+- template 配置：`ATCTrack-VLM/experiments/atctrack/atctrack_qwen_state.yaml:17`
+- search 配置：`ATCTrack-VLM/experiments/atctrack/atctrack_qwen_state.yaml:7`
+
+采样逻辑位置：`ATCTrack-VLM/lib/train/data/sampler.py:131`
+
+```text
+base_frame_id 作为第一个 template
+prev_frame_ids 随机取前面的 num_template_frames - 1 帧
+template_frame_ids = base_frame_id + prev_frame_ids
+search_frame_ids 从 base_frame_id 之后随机取 4 帧
+```
+
+tracker 使用逻辑位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:676`
+
+```text
+template[0]：tracker 固定使用的初始模板
+template[-2]：旧动态模板，给 Qwen 的 Frame 1
+template[-1]：新候选模板，给 Qwen 的 Frame 2
+
+Qwen 输出 yes：
+    dynamic_template = template[-1]
+Qwen 输出 no：
+    dynamic_template = template[-2]
+```
+
+具体选择位置：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:681`
+
+## 5. Loss 设计
+
+loss 位置：`ATCTrack-VLM/lib/train/actors/atctrack.py:101`
+
+主 loss：
+
+```text
+2.0 * giou_loss
++ 5.0 * l1_loss
++ focal_loss
++ confidence_loss
++ 0.01 * qwen_format_loss
++ 0.1 * qwen_teacher_loss
+```
+
+代码位置：`ATCTrack-VLM/lib/train/actors/atctrack.py:151`
+
+当前 `TRAIN.TYPE == target_state` 时，原来的 subject index loss 权重被置 0：`ATCTrack-VLM/lib/train/actors/atctrack.py:139`
+
+`qwen_format_loss` 是对 Qwen 生成 token 的 CE：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:383`
+
+`qwen_teacher_loss` 是 teacher yes/no 监督 loss，teacher 查询逻辑在：`ATCTrack-VLM/lib/models/atctrack/atctrack.py:247`
+
+注意：yaml 里 teacher 默认 false，但脚本里默认 true。
+
+## 6. 可训练参数
+
+optimizer 冻结/解冻位置：`ATCTrack-VLM/lib/train/base_functions.py:233`
+
+当前 `target_state` 模式默认可训练：
+
+```text
+target_state_encoder.projector
+target_state_encoder.film
+target_state_encoder.film_gate
+Qwen LoRA 权重，名字包含 lora_
+box_head
+confidence_pred
+Qwen input embedding 表，但 hook 只训练 <TARGET_STATE> 一行
+```
+
+embedding 单独 optimizer group，`weight_decay=0`：`ATCTrack-VLM/lib/train/base_functions.py:271`
+
+## 7. 训练配置
+
+配置文件：`ATCTrack-VLM/experiments/atctrack/atctrack_qwen_state.yaml:1`
+
+关键参数：
+
+```yaml
+MODEL.TARGET_STATE.ENABLE: true
+MODEL.TARGET_STATE.MODEL_PATH: "../Qwen/Qwen3.5-2B"
+MODEL.TARGET_STATE.FREEZE_QWEN: true
+MODEL.TARGET_STATE.TRAIN_TOKEN_EMBEDDING: true
+MODEL.TARGET_STATE.USE_LORA: true
+LORA_R: 8
+LORA_ALPHA: 16
+LORA_DROPOUT: 0.05
+
+TRAIN.BATCH_SIZE: 2
+TRAIN.EPOCH: 180
+TRAIN.LR: 0.0001
+TRAIN.WEIGHT_DECAY: 0.0001
+TRAIN.AMP: False
+TRAIN.TYPE: target_state
+```
+
+训练脚本：`ATCTrack-VLM/scripts/train_qwen_state.sh:1`
+
+路径参数位置：`ATCTrack-VLM/scripts/train_qwen_state.sh:8`
+
+```bash
+GPUS=0,1,2,3
+NPROC_PER_NODE=4
+QWEN_MODEL_PATH=../Qwen/Qwen3.5-2B
+ROBERTA_MODEL_PATH=resource/pretrained_models/roberta-base
+LASOT_DIR=../Dataset/LaSOT/LaSOTBenchmark
+```
+
+启动命令位置：`ATCTrack-VLM/scripts/train_qwen_state.sh:32`
+
+## 8. 保存方式
+
+轻量保存逻辑位置：`ATCTrack-VLM/lib/train/trainers/base_trainer.py:115`
+
+`target_state` 模式保存：
+
+```text
+box_head.*
+confidence_pred.*
+target_state_encoder.projector.*
+target_state_encoder.film.*
+target_state_encoder.film_gate
+所有 lora_ 权重
+<TARGET_STATE> embedding 那一行
+epoch / stats / settings
+```
+
+不保存：
+
+```text
+完整 Qwen base
+完整 Qwen embedding 表
+完整 backbone
+完整 RoBERTa
+optimizer state
+```
+
+具体保存 state 位置：`ATCTrack-VLM/lib/train/trainers/base_trainer.py:171`
+
+加载轻量 checkpoint 位置：`ATCTrack-VLM/lib/train/trainers/base_trainer.py:268`
+
+加载时会 `strict=False` 加载增量权重，然后单独恢复 `<TARGET_STATE>` embedding 行。
+
+默认保存路径：
+
+```text
+/media/data/WWZ/SX/ATCTrack-VLM/output/checkpoints/train/atctrack/atctrack_qwen_state/ATCTrack_epXXXX.pth.tar
+```
+
+保存频率：
+
+```text
+每 5 个 epoch 保存一次
+最后 10 个 epoch 每个 epoch 保存
+rank 0 保存
+```

README.md

@@ -0,0 +1,167 @@
+# ATCTrack: Aligning Target-Context Cues with Dynamic Target States for Robust Vision-Language Tracking
+
+> [Xiaokun Feng](https://scholar.google.com.hk/citations?user=NqXtIPIAAAAJ),  [Shiyu Hu](https://huuuuusy.github.io/), [Xuchen Li](https://github.com/Xuchen-Li), [Dailing Zhang](https://scholar.google.com.hk/citations?user=ApH4wOcAAAAJ), [Meiqi Wu](https://scholar.google.com.hk/citations?user=fGc7NVAAAAAJ), [Jing Zhang](https://github.com/XiaokunFeng/CSTrack), [Xiaotang Chen](http://www.ia.cas.cn/rcdw/fyjy/202404/t20240422_7129814.html), [Kaiqi Huang](https://people.ucas.ac.cn/~0004554)
+
+
+[![](https://img.shields.io/badge/ATCTrack-arXiv%20-b31b1b?logo=arxiv&logoColor=red)](https://arxiv.org/abs/2507.19875)
+[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20ATCTrack-Results&ckpts-red)](https://huggingface.co/Xiaokunfeng2022/ATCTrack)
+
+This is an official pytorch implementation of the paper **ATCTrack: Aligning Target-Context Cues with Dynamic Target States for Robust Vision-Language Tracking**.
+
+
+### 🔥 Updates
+
+*   \[8/2025\] **ATCTrack's** code is available!
+*   \[6/2025\] **ATCTrack**  is accepted by ICCV25 Highlight!
+
+### 📣 Overview
+#### Our motivation & Core modeling approach
+
+Vision-language tracking aims to locate the target object in the video sequence using a template patch and a language description provided in the initial frame. To achieve
+robust tracking, especially in complex long-term scenarios that reflect real-world conditions as recently highlighted by
+MGIT, it is essential not only to characterize the target features but also to utilize the context features related to the
+target. However, the visual and textual target-context cues
+derived from the initial prompts generally align only with
+the initial target state. Due to their dynamic nature, target states are constantly changing, particularly in complex
+long-term sequences. It is intractable for these cues to continuously guide Vision-Language Trackers (VLTs). Furthermore, for the text prompts with diverse expressions, our
+experiments reveal that existing VLTs struggle to discern
+which words pertain to the target or the context, complicating the utilization of textual cues. 
+![ATCTrack_motivation](asset/motivation.png)
+
+In this work, we present a novel tracker named ATCTrack, which can obtain multimodal cues Aligned with the dynamic target states
+through comprehensive Target-Context feature modeling,
+thereby achieving robust tracking. Specifically, (1) for the
+visual modality, we propose an effective temporal visual
+target-context modeling approach that provides the tracker
+with timely visual cues. (2) For the textual modality, we
+achieve precise target words identification solely based on
+textual content, and design an innovative context words
+calibration method to adaptively utilize auxiliary context
+words. (3) We conduct extensive experiments on mainstream benchmarks and ATCTrack achieves a new SOTA
+performance
+![ATCTrack_pipeline](asset/framework.png)
+
+#### Strong performance
+
+![ATCTrack_experiment](asset/experiment.png)
+
+
+
+### 🔨 Installation
+```
+conda create -n atctrack python=3.8
+conda activate atctrack
+bash install.sh
+```
+
+### 🔧 Usage
+
+#### Data Preparation
+Our ATCTrack is trained on  LaSOT, TNL2K, RefCOCOg, OTB99-Lang, VastTrack, GOT-10k, and TrackingNet datasets.  
+Put these tracking datasets in [./data](data). It should look like:
+
+   ```
+   ${ATCTrack_ROOT}
+    -- data
+        -- lasot
+            |-- airplane
+            |-- basketball
+            |-- bear
+            ...
+        -- got10k
+            |-- test
+            |-- train
+            |-- val
+        -- coco
+            |-- annotations
+            |-- images
+        -- trackingnet
+            |-- TRAIN_0
+            |-- TRAIN_1
+            ...
+            |-- TRAIN_11
+            |-- TEST
+        -- VastTrack
+            |-- unisot_train_final_backup
+                |-- Aardwolf
+                ...
+                |-- Zither
+            |-- unisot_final_test
+                |-- Aardwolf
+                ...
+                |-- Zither
+        -- tnl2k
+            -- train
+                |-- Arrow_Video_ZZ04_done
+                |-- Assassin_video_1-Done
+                |-- Assassin_video_2-Done
+                ...
+            -- test
+                |-- advSamp_Baseball_game_002-Done
+                |-- advSamp_Baseball_video_01-Done
+                |-- advSamp_Baseball_video_02-Done
+                ...
+
+   ```
+
+#### Set project paths
+Run the following command to set paths for this project
+```
+python tracking/create_default_local_file.py --workspace_dir . --data_dir ./data --save_dir .
+```
+After running this command, you can also modify paths by editing these two files
+```
+lib/train/admin/local.py  # paths about training
+lib/test/evaluation/local.py  # paths about testing
+```
+
+#### Train
+##### Prepare pretrained backbone
+The backbone and patch embedding of ATCTrack are initialized with pre-trained weights from [**Fast-iTPN**](https://github.com/sunsmarterjie/iTPN), and we adopt RoBERTa-Base as our text encoder.  
+Please download the **fast_itpn_base_clipl_e1600.pt**, **fast_itpn_large_1600e_1k.pt** and **roberta-base** checkpoints and place them in [./resource/pretrained_models](./resource/pretrained_models).
+
+##### Train ATCTrack
+You can run the following command to train the ATCTrack-B:
+```
+python -m torch.distributed.launch --nproc_per_node 3 lib/train/run_training.py --script atctrack --config atctrack_base --save_dir .
+```
+
+Besides, you can run the following command to train the ATCTrack-L:
+```
+python -m torch.distributed.launch --nproc_per_node 3 lib/train/run_training.py --script atctrack --config atctrack_large --save_dir .
+```
+
+#### Test and evaluate on benchmarks
+First, you need to set the paths for the various evaluation benchmarks in [./lib/test/evaluation/local.py](./lib/test/evaluation/local.py), and prepare the model weights for evaluation. 
+Then, run the following command to perform evaluation on different benchmarks (taking atctrack_base as an example).
+- LaSOT
+```
+python tracking/test.py --tracker_name atctrack --tracker_param atctrack_base --dataset lasot_lang --threads 4 --num_gpus 2 --ckpt_path '{your_dir_saved_model_ckpt}/ATCTrack_b.pth.tar'
+python tracking/analysis_results.py --dataset_name lasot_lang --tracker_param atctrack_base
+```
+- LaSOT_ext
+```
+python tracking/test.py --tracker_name atctrack --tracker_param atctrack_base --dataset lasot_extension_subset_lang --threads 4 --num_gpus 2 --ckpt_path '{your_dir_saved_model_ckpt}/ATCTrack_b.pth.tar'
+python tracking/analysis_results.py --dataset_name lasot_extension_subset_lang --tracker_param atctrack_base
+```
+
+- TNL2K
+```
+python tracking/test.py --tracker_name atctrack --tracker_param atctrack_base --dataset tnl2k --threads 4 --num_gpus 2 --ckpt_path '{your_dir_saved_model_ckpt}/ATCTrack_b.pth.tar'
+python tracking/analysis_results.py --dataset_name lasot_extension_subset_lang --tracker_param atctrack_base
+```
+
+- MGIT
+
+Please refer to the [official MGIT testing platform](http://videocube.aitestunion.com/) and [tools](https://github.com/huuuuusy/videocube-toolkit) to complete the corresponding evaluation.
+
+### 📊 Model Zoo
+The trained models, and the raw tracking results are provided in the [![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20ATCTrack-Results&ckpts-red)](https://huggingface.co/Xiaokunfeng2022/ATCTrack).
+
+
+### ❤️Acknowledgement
+We would like to express our gratitude to the following open-source repositories that our work is based on: [SeqtrackV2](https://github.com/chenxin-dlut/SeqTrackv2),  [AQATrack](https://github.com/GXNU-ZhongLab/AQATrack), [Fast-iTPN](https://github.com/sunsmarterjie/iTPN).
+Their contributions have been invaluable to this project.
+
+
+

TRAIN_QWEN_STATE_GUIDE.md

@@ -0,0 +1,405 @@
+# ATCTrack-VLM Qwen Target-State 训练指南
+
+本文档说明如何在当前 ATCTrack-VLM 代码中运行 Qwen target-state 更新器训练，以及迁移到其他服务器时需要检查的路径、参数和保存方式。内容已按当前代码实现同步。
+
+## 1. 当前训练做什么
+
+当前方法以 ATCTrack 为 baseline，新增一个基于 Qwen 的 target-state 更新器。
+
+和原始 ATCTrack 的主要区别：
+
+- 每个训练样本使用 3 张 template 和 4 张 search。
+- tracker 始终使用 `template[0]` 作为静态模板。
+- Qwen 比较 `template[-2]` 和 `template[-1]`，两张图都会裁剪到 template size，并画上 bbox。
+- Qwen 输入原始 prompt 和两张图，训练/推理时都只 `generate` 一次。
+- Qwen 输出被约束为两种合法格式之一：
+
+```text
+<answer>yes</answer><state_token><TARGET_STATE></state_token>
+<answer>no</answer><state_token><TARGET_STATE></state_token>
+```
+
+- 如果 Qwen 输出 `yes`，tracker 使用 `template[-1]` 作为动态模板。
+- 如果 Qwen 输出 `no`，tracker 使用 `template[-2]` 作为动态模板。
+- 提取 `<TARGET_STATE>` 对应 hidden state，经过 projector 映射到 tracker 维度，再通过 FiLM 调制 tracker feature。
+- 可选 teacher 模型输出 `<answer>yes/no</answer>`，用于监督当前 Qwen 更新器的 yes/no 判断。
+
+## 2. 推荐目录结构
+
+默认脚本假设目录结构如下：
+
+```text
+parent_dir/
+  ATCTrack-VLM/
+  Qwen/
+    Qwen3.5-2B/
+  Dataset/
+    LaSOT/
+      LaSOTBenchmark/
+```
+
+如果目录不同，可以用环境变量覆盖路径。
+
+必需文件：
+
+- ATCTrack checkpoint：`ATCTrack-VLM/checkpoint/ATCTrack_b.pth.tar`
+- backbone 预训练文件：`ATCTrack-VLM/resource/pretrained_models/fast_itpn_base_clipl_e1600.pt`
+- RoBERTa：`ATCTrack-VLM/resource/pretrained_models/roberta-base`
+- Qwen 更新器：默认 `ATCTrack-VLM/../Qwen/Qwen3.5-2B`
+- LaSOT：默认 `ATCTrack-VLM/../Dataset/LaSOT/LaSOTBenchmark`
+
+## 3. 训练脚本和环境变量
+
+主训练脚本：
+
+```bash
+scripts/train_qwen_state.sh
+```
+
+脚本会自动根据自身位置推导 `PROJECT_DIR`，所以项目移动到其他服务器后通常不需要改脚本，只需要保证相对目录或环境变量正确。
+
+当前脚本里的关键变量：
+
+```bash
+PROJECT_DIR=/path/to/ATCTrack-VLM
+PYTHON_BIN=python
+GPUS=0,1,2,3
+NPROC_PER_NODE=4
+MASTER_PORT=29501
+SAVE_DIR=/path/to/output
+SEED=0
+CONFIG=atctrack_qwen_state
+
+QWEN_MODEL_PATH=/path/to/Qwen/Qwen3.5-2B
+ROBERTA_MODEL_PATH=/path/to/ATCTrack-VLM/resource/pretrained_models/roberta-base
+LASOT_DIR=/path/to/LaSOT/LaSOTBenchmark
+
+QWEN_TEACHER_ENABLE=true
+QWEN_TEACHER_MODEL=qwen3.5
+QWEN_TEACHER_BASE_URL=http://127.0.0.1:8001/v1
+QWEN_TEACHER_API_KEY=sk-no-key-required
+```
+
+注意：虽然 yaml 里 `MODEL.TARGET_STATE.TEACHER_ENABLE: false`，但脚本默认：
+
+```bash
+QWEN_TEACHER_ENABLE=true
+```
+
+代码优先读取环境变量，所以直接运行脚本会调用 teacher。如果不想调用 teacher，使用：
+
+```bash
+QWEN_TEACHER_ENABLE=false bash scripts/train_qwen_state.sh
+```
+
+数据集路径也可以用环境变量覆盖：
+
+```bash
+DATASET_ROOT=/path/to/Dataset
+LASOT_DIR=/path/to/LaSOT/LaSOTBenchmark
+```
+
+## 4. Teacher 服务
+
+如果 `QWEN_TEACHER_ENABLE=true`，训练会调用 OpenAI-compatible API 作为 teacher，生成 yes/no 标签。
+
+teacher prompt：
+
+```text
+Role: [OBJECT] tracking update judge.
+Task: Compare the targets inside the provided bboxes in Frame 1 (Original) and Frame 2 (New), and decide whether Frame 2 should update the tracking template.
+
+Reject update for full occlusion, out of view, too small target, severe blur/clipping, wrong bbox, distractor, uncertain identity, or no meaningful target appearance change.
+Accept only if Frame 2 is the same target, bbox is reliable, target is clear, and appearance change is useful.
+
+Output exactly: <answer>yes/no</answer>
+```
+
+vLLM 服务示例：
+
+```bash
+python -m vllm.entrypoints.openai.api_server \
+  --model /path/to/Qwen/Qwen3.5-2B \
+  --served-model-name qwen3.5 \
+  --port 8001 \
+  --tensor-parallel-size 4 \
+  --trust-remote-code \
+  --dtype bfloat16 \
+  --gpu-memory-utilization 0.4 \
+  --max-model-len 40000 \
+  --enforce-eager
+```
+
+建议不要把 teacher 服务和 4-rank DDP 训练放在同一组 24GB GPU 上。之前测试中 teacher 服务和训练共享 GPU 时显存非常接近 3090 上限。
+
+如果没有 teacher 服务，先关闭 teacher：
+
+```bash
+QWEN_TEACHER_ENABLE=false bash scripts/train_qwen_state.sh
+```
+
+## 5. 启动训练
+
+先激活环境：
+
+```bash
+conda activate VLT
+cd /path/to/ATCTrack-VLM
+bash scripts/train_qwen_state.sh
+```
+
+迁移到其他服务器时，推荐显式覆盖路径：
+
+```bash
+PROJECT_DIR=/path/to/ATCTrack-VLM \
+QWEN_MODEL_PATH=/path/to/Qwen/Qwen3.5-2B \
+ROBERTA_MODEL_PATH=/path/to/ATCTrack-VLM/resource/pretrained_models/roberta-base \
+LASOT_DIR=/path/to/LaSOT/LaSOTBenchmark \
+SAVE_DIR=/path/to/output \
+GPUS=0,1,2,3 \
+NPROC_PER_NODE=4 \
+QWEN_TEACHER_ENABLE=false \
+bash scripts/train_qwen_state.sh
+```
+
+单卡 smoke test：
+
+```bash
+GPUS=0 NPROC_PER_NODE=1 QWEN_TEACHER_ENABLE=false bash scripts/train_qwen_state.sh
+```
+
+## 6. 当前可训练和冻结模块
+
+当前 `TRAIN.TYPE: target_state` 时，可训练部分为：
+
+- ATCTrack `box_head`
+- ATCTrack `confidence_pred`
+- Qwen LoRA 权重，参数名中包含 `lora_`
+- `target_state_encoder.projector`
+- `target_state_encoder.film`
+- `target_state_encoder.film_gate`
+- Qwen 原始 input embedding 表里的 `<TARGET_STATE>` 那一行
+
+`<TARGET_STATE>` embedding 不是独立外挂参数。当前做法是打开 Qwen 原始 embedding 表的梯度，然后通过 grad hook 只保留 `<TARGET_STATE>` 对应 token id 那一行梯度，其它 token embedding 梯度置 0。
+
+embedding 参数被单独放到 optimizer group，并设置：
+
+```text
+weight_decay = 0.0
+```
+
+这样可以避免 AdamW 的 decoupled weight decay 改动整张 embedding 表。
+
+冻结部分：
+
+- Qwen base 参数，除了 LoRA 和 `<TARGET_STATE>` embedding 行。
+- ATCTrack backbone/encoder。
+- RoBERTa text encoder。
+- 其它不在 target_state 可训练关键词里的模块。
+
+## 7. Loss 设计
+
+主 loss：
+
+```text
+L = 2.0 * GIoU
+  + 5.0 * L1
+  + focal/location
+  + confidence
+  + QWEN_FORMAT_WEIGHT * qwen_format_loss
+  + QWEN_TEACHER_WEIGHT * qwen_teacher_loss
+```
+
+当前配置：
+
+```yaml
+QWEN_FORMAT_WEIGHT: 0.01
+QWEN_TEACHER_WEIGHT: 0.1
+```
+
+说明：
+
+- `qwen_format_loss` 是基于 Qwen `generate` 返回的 scores，对实际生成的固定格式 token 计算 CE。由于输出受到 `prefix_allowed_tokens_fn` 约束，它主要是格式生成置信度约束，不代表格式一定错。
+- `qwen_teacher_loss` 使用 teacher yes/no 标签监督当前 Qwen 在 yes/no token 位置的判断。
+- 当 teacher 关闭或 teacher 返回无效时，`qwen_teacher_loss = 0`。
+- `TRAIN.TYPE == target_state` 时，原 subject index loss 权重为 0。
+
+## 8. 当前训练配置
+
+配置文件：
+
+```text
+experiments/atctrack/atctrack_qwen_state.yaml
+```
+
+关键配置：
+
+```yaml
+DATA:
+  TEMPLATE:
+    NUMBER: 3
+    SIZE: 128
+  SEARCH:
+    NUMBER: 4
+    SIZE: 256
+  TRAIN:
+    DATASETS_NAME:
+    - LASOT
+    DATASETS_RATIO:
+    - 1
+    SAMPLE_PER_EPOCH: 25000
+
+MODEL:
+  PRETRAIN_FILE: "fast_itpn_base_clipl_e1600.pt"
+  PRETRAINED_PATH: "checkpoint/ATCTrack_b.pth.tar"
+  TARGET_STATE:
+    ENABLE: true
+    MODEL_PATH: "../Qwen/Qwen3.5-2B"
+    TOKEN: "<TARGET_STATE>"
+    FREEZE_QWEN: true
+    TRAIN_TOKEN_EMBEDDING: true
+    USE_LORA: true
+    LORA_R: 8
+    LORA_ALPHA: 16
+    LORA_DROPOUT: 0.05
+    TEACHER_ENABLE: false
+
+TRAIN:
+  BATCH_SIZE: 2
+  EPOCH: 180
+  LR: 0.0001
+  LR_DROP_EPOCH: 144
+  WEIGHT_DECAY: 0.0001
+  GRAD_CLIP_NORM: 0.1
+  AMP: False
+  TYPE: target_state
+```
+
+注意：teacher 是否开启以环境变量 `QWEN_TEACHER_ENABLE` 为准。脚本默认是 true。
+
+## 9. 显存记录
+
+最近一次 2B、单进程、batch size 2、关闭 teacher 的完整 forward/backward/optimizer step 测试结果：
+
+```text
+结果目录：temp-result/qwen35_2b_train_memory_after_target_embedding_20260608_121711
+
+after_model_to_cuda: 5549 MiB
+after_forward_loss: 13405 MiB
+after_backward: 16357 MiB
+after_optimizer_step: 17329 MiB
+
+torch max allocated: 15572.61 MiB
+```
+
+因此在 3090 24GB 上，2B + batch size 2 + 无 teacher 的单进程训练大约需要 16 到 17GB 常驻显存。
+
+如果启用 teacher，尤其 teacher 服务也占用同一张 GPU，显存会明显增加。DDP 训练时每个 rank 都会有一份训练模型，并且每个 rank 都可能调用 teacher API。
+
+建议：
+
+- 先用 `QWEN_TEACHER_ENABLE=false` 做 smoke test。
+- teacher 服务尽量放到单独 GPU 或单独机器。
+- 9B 不适合普通 DDP 复制到 24GB GPU 上训练，除非使用模型并行、FSDP、DeepSpeed 或外部推理服务。
+
+## 10. 轻量保存和恢复
+
+当前 target_state 模式使用轻量 checkpoint 保存。
+
+保存位置：
+
+```text
+output/checkpoints/train/atctrack/atctrack_qwen_state/ATCTrack_epXXXX.pth.tar
+```
+
+保存频率：
+
+```text
+每 5 个 epoch 保存一次
+最后 10 个 epoch 每个 epoch 保存
+rank 0 保存
+```
+
+轻量 checkpoint 保存内容：
+
+```text
+box_head.*
+confidence_pred.*
+target_state_encoder.projector.*
+target_state_encoder.film.*
+target_state_encoder.film_gate
+所有 lora_ 权重
+<TARGET_STATE> embedding 那一行
+epoch / stats / settings
+```
+
+不保存：
+
+```text
+完整 Qwen base
+完整 Qwen embedding 表
+完整 backbone
+完整 RoBERTa
+optimizer state
+```
+
+可以加载轻量 checkpoint 继续训练。加载时会：
+
+```text
+1. strict=False 加载轻量 net 权重
+2. 单独恢复 <TARGET_STATE> embedding 那一行
+3. 恢复 epoch 和 stats
+```
+
+注意：轻量 checkpoint 不保存 optimizer state，所以继续训练时 AdamW 动量不会恢复。这是为了减小 checkpoint 体积的取舍。
+
+## 11. 迁移路径检查
+
+生产训练路径���经做了相对路径和环境变量兼容：
+
+- `scripts/train_qwen_state.sh` 会根据脚本位置推导 `PROJECT_DIR`。
+- `QWEN_MODEL_PATH`、`ROBERTA_MODEL_PATH`、`LASOT_DIR`、teacher URL/model/key 都可以用环境变量覆盖。
+- `lib/train/admin/local.py` 使用 `DATASET_ROOT` 和各数据集环境变量。
+- `lib/test/evaluation/local.py` 也使用环境变量。
+
+之前扫描到的绝对路径主要在非当前训练入口文件里：
+
+- `lib/train/run_training_debug_version.py`
+- `lib/utils/gen_fname_list.py`
+- `lib/utils/write_to_lmdb.py`
+- VideoCube 相关 test 代码
+- `temp-result` 下的调试脚本
+
+这些不会被 `scripts/train_qwen_state.sh` 调用。只有使用这些工具或 VideoCube 评测时才需要另外修改。
+
+## 12. 新服务器训练前检查清单
+
+1. 检查 Python 环境：
+
+```bash
+python -c "import torch, transformers, peft, openai; print(torch.__version__)"
+```
+
+2. 检查关键文件：
+
+```bash
+ls checkpoint/ATCTrack_b.pth.tar
+ls resource/pretrained_models/fast_itpn_base_clipl_e1600.pt
+ls resource/pretrained_models/roberta-base
+ls "$QWEN_MODEL_PATH"
+ls "$LASOT_DIR"
+```
+
+3. 如果开启 teacher，检查 API：
+
+```bash
+curl http://127.0.0.1:8001/v1/models
+```
+
+4. 先跑单卡 smoke test：
+
+```bash
+GPUS=0 NPROC_PER_NODE=1 QWEN_TEACHER_ENABLE=false bash scripts/train_qwen_state.sh
+```
+
+5. smoke test 正常后，再按目标 GPU 数量和 teacher 设置启动完整训练。

VLT_environment_summary.md

@@ -0,0 +1,223 @@
+# VLT 虚拟环境迁移说明
+
+本文档总结当前机器上的 `VLT` conda 环境，用于在另一台电脑上重新配置相近环境。
+
+## 1. 当前环境概况
+
+- 环境名：`VLT`
+- Python：`3.11.14`
+- PyTorch：`2.10.0+cu128`
+- torchvision：`0.25.0+cu128`
+- torchaudio：`2.10.0`
+- torchdata：`0.11.0`
+- PyTorch CUDA runtime：`12.8`
+- cuDNN：`9.10.2`
+- NCCL：`2.27.5`
+- 当前机器 NVIDIA Driver：`560.35.05`
+- 当前机器 CUDA Driver API：`12.6`
+
+注意：这个环境里的 PyTorch 是 CUDA 12.8 wheel。新机器的 NVIDIA 驱动需要足够新，建议使用 `560+` 或更新版本。
+
+## 2. 推荐创建方式
+
+建议先创建干净 conda 环境，再用 pip 安装核心依赖：
+
+```bash
+conda create -n VLT python=3.11.14 pip -y
+conda activate VLT
+```
+
+安装 PyTorch CUDA 12.8 版本：
+
+```bash
+pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cu128
+```
+
+然后安装视觉跟踪、VLM、训练相关依赖：
+
+```bash
+pip install \
+  timm==1.0.25 \
+  numpy==2.3.5 \
+  opencv-python==4.13.0.92 \
+  opencv-python-headless==4.13.0.92 \
+  pillow==11.3.0 \
+  scipy==1.17.0 \
+  scikit-image==0.26.0 \
+  pandas==2.3.3 \
+  matplotlib==3.10.8 \
+  pycocotools==2.0.11 \
+  lmdb==1.8.1 \
+  easydict==1.13 \
+  omegaconf==2.3.0 \
+  einops==0.8.2 \
+  tqdm==4.67.1 \
+  visdom==0.2.4 \
+  tensorboard==2.20.0 \
+  tensorboardX==2.6.4 \
+  nltk==3.9.4 \
+  openai==2.29.0 \
+  transformers==5.10.2 \
+  accelerate==1.11.0 \
+  datasets==4.0.0 \
+  peft==0.18.0 \
+  trl==0.24.0 \
+  deepspeed==0.18.4 \
+  gradio==5.50.0 \
+  fastapi==0.128.0 \
+  uvicorn==0.40.0 \
+  modelscope==1.34.0
+```
+
+安装 OpenMMLab 相关包：
+
+```bash
+pip install mmengine==0.10.7 mmdet==3.3.0 mmcv==2.1.0
+```
+
+如果 `mmcv` 安装失败，优先根据新机器的 PyTorch/CUDA 版本选择 OpenMMLab 官方预编译 wheel；不要盲目源码编译。
+
+## 3. 当前环境中的关键包版本
+
+### 跟踪/视觉相关
+
+```text
+timm                      1.0.25
+opencv-python             4.13.0.92
+opencv-python-headless    4.13.0.92
+Pillow                    11.3.0
+numpy                     2.3.5
+scipy                     1.17.0
+scikit-image              0.26.0
+pycocotools               2.0.11
+lmdb                      1.8.1
+easydict                  1.13
+omegaconf                 2.3.0
+einops                    0.8.2
+visdom                    0.2.4
+```
+
+### PyTorch/CUDA 相关
+
+```text
+torch                     2.10.0+cu128
+torchvision               0.25.0+cu128
+torchaudio                2.10.0
+torchdata                 0.11.0
+triton                    3.6.0
+nvidia-cublas-cu12        12.8.4.1
+nvidia-cudnn-cu12         9.10.2.21
+nvidia-nccl-cu12          2.27.5
+```
+
+### MLLM/VLM/API 相关
+
+```text
+openai                    2.29.0
+transformers              5.10.2
+tokenizers                0.22.2
+accelerate                1.11.0
+datasets                  4.0.0
+peft                      0.18.0
+trl                       0.24.0
+deepspeed                 0.18.4
+sentencepiece             0.2.1
+tiktoken                  0.12.0
+modelscope                1.34.0
+gradio                    5.50.0
+fastapi                   0.128.0
+uvicorn                   0.40.0
+```
+
+### OpenMMLab 相关
+
+```text
+mmcv                      2.1.0
+mmdet                     3.3.0
+mmengine                  0.10.7
+```
+
+## 4. 本地 editable 包
+
+当前环境里有一个本地 editable 安装：
+
+```text
+llamafactory 0.9.5.dev0 -> /media/data/WWZ/SX/LlamaFactory
+```
+
+如果新机器也需要跑 LLaMA-Factory 相关训练或推理，需要把该仓库同步过去后执行：
+
+```bash
+cd /path/to/LlamaFactory
+pip install -e .
+```
+
+如果只跑 ATCTrack/SUTrack/CTVLT 等跟踪测试，不一定需要这个 editable 包。
+
+## 5. NLTK 数据
+
+部分代码会下载：
+
+```text
+punkt
+averaged_perceptron_tagger
+```
+
+如果新机器无法联网，可以从旧机器的 `~/nltk_data` 拷贝到新机器同路径，或手动设置：
+
+```bash
+export NLTK_DATA=/path/to/nltk_data
+```
+
+## 6. 代理设置
+
+如果新机器访问 HuggingFace、PyPI 或 OpenAI-compatible API 需要代理，可以写入 shell 配置：
+
+```bash
+export HTTP_PROXY=http://127.0.0.1:7890
+export HTTPS_PROXY=http://127.0.0.1:7890
+export ALL_PROXY=http://127.0.0.1:7890
+```
+
+## 7. 验证命令
+
+环境配置完成后，先验证 PyTorch/CUDA：
+
+```bash
+python - <<'PY'
+import torch, torchvision
+print("torch:", torch.__version__)
+print("torchvision:", torchvision.__version__)
+print("cuda available:", torch.cuda.is_available())
+print("torch cuda:", torch.version.cuda)
+print("cudnn:", torch.backends.cudnn.version())
+if torch.cuda.is_available():
+    print("gpu:", torch.cuda.get_device_name(0))
+PY
+```
+
+验证跟踪常用依赖：
+
+```bash
+python - <<'PY'
+import cv2, timm, numpy, pandas, mmcv, mmdet, mmengine
+import transformers, openai
+print("basic imports ok")
+PY
+```
+
+## 8. 复现���整环境的备选方式
+
+如果希望尽量完整复刻当前环境，可以在旧机器导出：
+
+```bash
+conda env export -n VLT --no-builds > vlt_environment.yml
+```
+
+在新机器创建：
+
+```bash
+conda env create -f vlt_environment.yml
+```
+
+但这种方式可能因为 CUDA wheel、OpenMMLab wheel、镜像源或本地 editable 包路径不同而失败。实际迁移时更推荐按本文第 2 节分组安装。

checkpoint/ATCTrack_b.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c01b720243b03216554dd8c41bcd9f14136cee14847929bcfe94a32b6aa7f00e
+size 1703949574

experiments/atctrack/atctrack_base.yaml

@@ -0,0 +1,103 @@
+DATA:
+  MAX_SAMPLE_INTERVAL: 400
+  MEAN:
+  - 0.485
+  - 0.456
+  - 0.406
+  SEARCH:
+    CENTER_JITTER: 3.5
+    FACTOR: 4.0
+    SCALE_JITTER: 0.5
+    SIZE: 256
+    NUMBER: 4 # 6
+  STD:
+  - 0.229
+  - 0.224
+  - 0.225
+  TEMPLATE:
+    CENTER_JITTER: 0
+    FACTOR: 2.0
+    SCALE_JITTER: 0
+    SIZE: 128
+    NUMBER: 2
+  TRAIN:
+    DATASETS_NAME:
+    - LASOT
+    # - GOT10K_vottrain
+    # - TRACKINGNET
+    # - VastTrack
+    # - TNL2K_train
+    # - RefCOCO14    # RefCOCO14,COCO17
+    # - OTB99_train
+    DATASETS_RATIO:
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    - 1
+    SAMPLE_PER_EPOCH: 25000
+MODEL:
+  PRETRAIN_FILE: "fast_itpn_base_clipl_e1600.pt"    # for backbone
+  PRETRAINED_PATH: "checkpoint/ATCTrack_b.pth.tar"
+  TARGET_STATE:
+    ENABLE: true
+    MODEL_PATH: "../Qwen/Qwen3.5-4B"
+    TOKEN: "<TARGET_STATE>"
+    FREEZE_QWEN: true
+    TRAIN_TOKEN_EMBEDDING: true
+    USE_LORA: true
+    LORA_R: 8
+    LORA_ALPHA: 16
+    LORA_DROPOUT: 0.05
+    LORA_TARGET_MODULES:
+      - in_proj_qkv
+      - out_proj
+      - in_proj_z
+      - in_proj_b
+      - in_proj_a
+      - q_proj
+      - k_proj
+      - v_proj
+      - o_proj
+      - gate_proj
+      - up_proj
+      - down_proj
+  BACKBONE:
+    TYPE: itpn_base
+    STRIDE: 16
+  HEAD:
+    TYPE: CENTER
+    NUM_CHANNELS: 256
+  HIDDEN_DIM: 512
+TRAIN:
+  BACKBONE_MULTIPLIER: 0.1
+  DROP_PATH_RATE: 0.1
+  BATCH_SIZE: 1
+  EPOCH: 180
+  GIOU_WEIGHT: 2.0
+  L1_WEIGHT: 5.0
+  GRAD_CLIP_NORM: 0.1
+  LR: 0.0001
+  LR_DROP_EPOCH: 144  # 4/5（0.8）
+  NUM_WORKER: 8
+  OPTIMIZER: ADAMW
+  PRINT_INTERVAL: 50
+  SCHEDULER:
+    TYPE: step
+    DECAY_RATE: 0.1
+  WEIGHT_DECAY: 0.0001
+  AMP: False
+  FIX_BN: true
+  TYPE: target_state # train target-state projector/gate and tracker head
+
+TEST:
+  EPOCH: 240
+  SEARCH_FACTOR: 4.0
+  SEARCH_SIZE: 256
+  TEMPLATE_FACTOR: 2.0
+  TEMPLATE_SIZE: 128
+  WINDOW: true
+  NUM_TEMPLATES: 2
+

experiments/atctrack/atctrack_large.yaml

@@ -0,0 +1,79 @@
+DATA:
+  MAX_SAMPLE_INTERVAL: 400
+  MEAN:
+  - 0.485
+  - 0.456
+  - 0.406
+  SEARCH:
+    CENTER_JITTER: 3.5
+    FACTOR: 4.0
+    SCALE_JITTER: 0.5
+    SIZE: 256
+    NUMBER: 4 # 6
+  STD:
+  - 0.229
+  - 0.224
+  - 0.225
+  TEMPLATE:
+    CENTER_JITTER: 0
+    FACTOR: 2.0
+    SCALE_JITTER: 0
+    SIZE: 128
+    NUMBER: 2
+  TRAIN:
+    DATASETS_NAME:
+    - LASOT
+    - GOT10K_vottrain
+    - TRACKINGNET
+    - VastTrack
+    - TNL2K_train
+    - RefCOCO14    # RefCOCO14,COCO17
+    - OTB99_train
+    DATASETS_RATIO:
+    - 6
+    - 6
+    - 6
+    - 6
+    - 6
+    - 6
+    - 1
+    SAMPLE_PER_EPOCH: 25000
+MODEL:
+  PRETRAIN_FILE: "fast_itpn_large_1600e_1k.pt" # "fast_itpn_base_clipl_e1600.pt"    # for backbone
+  PRETRAINED_PATH: ""
+  BACKBONE:
+    TYPE: itpn_large
+    STRIDE: 16
+  HEAD:
+    TYPE: CENTER
+    NUM_CHANNELS: 256
+  HIDDEN_DIM: 512
+TRAIN:
+  BACKBONE_MULTIPLIER: 0.1
+  DROP_PATH_RATE: 0.1
+  BATCH_SIZE: 5
+  EPOCH: 150
+  GIOU_WEIGHT: 2.0
+  L1_WEIGHT: 5.0
+  GRAD_CLIP_NORM: 0.1
+  LR: 0.0001
+  LR_DROP_EPOCH: 120  # 4/5（0.8）
+  NUM_WORKER: 8
+  OPTIMIZER: ADAMW
+  PRINT_INTERVAL: 50
+  SCHEDULER:
+    TYPE: step
+    DECAY_RATE: 0.1
+  WEIGHT_DECAY: 0.0001
+  AMP: False
+  FIX_BN: true
+
+TEST:
+  EPOCH: 240
+  SEARCH_FACTOR: 4.0
+  SEARCH_SIZE: 256
+  TEMPLATE_FACTOR: 2.0
+  TEMPLATE_SIZE: 128
+  WINDOW: true
+  NUM_TEMPLATES: 2
+

experiments/atctrack/atctrack_qwen_state.yaml

@@ -0,0 +1,109 @@
+DATA:
+  MAX_SAMPLE_INTERVAL: 400
+  MEAN:
+  - 0.485
+  - 0.456
+  - 0.406
+  SEARCH:
+    CENTER_JITTER: 3.5
+    FACTOR: 4.0
+    SCALE_JITTER: 0.5
+    SIZE: 256
+    NUMBER: 4 # 6
+  STD:
+  - 0.229
+  - 0.224
+  - 0.225
+  TEMPLATE:
+    CENTER_JITTER: 0
+    FACTOR: 2.0
+    SCALE_JITTER: 0
+    SIZE: 128
+    NUMBER: 3
+  TRAIN:
+    DATASETS_NAME:
+    - LASOT
+    # - GOT10K_vottrain
+    # - TRACKINGNET
+    # - VastTrack
+    # - TNL2K_train
+    # - RefCOCO14    # RefCOCO14,COCO17
+    # - OTB99_train
+    DATASETS_RATIO:
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    # - 6
+    - 1
+    SAMPLE_PER_EPOCH: 25000
+MODEL:
+  PRETRAIN_FILE: "fast_itpn_base_clipl_e1600.pt"    # for backbone
+  PRETRAINED_PATH: "checkpoint/ATCTrack_b.pth.tar"
+  TARGET_STATE:
+    ENABLE: true
+    MODEL_PATH: "../Qwen/Qwen3.5-2B"
+    TOKEN: "<TARGET_STATE>"
+    FREEZE_QWEN: true
+    TRAIN_TOKEN_EMBEDDING: true
+    USE_LORA: true
+    LORA_R: 8
+    LORA_ALPHA: 16
+    LORA_DROPOUT: 0.05
+    LORA_TARGET_MODULES:
+      - in_proj_qkv
+      - out_proj
+      - in_proj_z
+      - in_proj_b
+      - in_proj_a
+      - q_proj
+      - k_proj
+      - v_proj
+      - o_proj
+      - gate_proj
+      - up_proj
+      - down_proj
+    TEACHER_ENABLE: false
+    TEACHER_MODEL: qwen3.5
+    TEACHER_BASE_URL: http://127.0.0.1:8001/v1
+    TEACHER_API_KEY: sk-no-key-required
+  BACKBONE:
+    TYPE: itpn_base
+    STRIDE: 16
+  HEAD:
+    TYPE: CENTER
+    NUM_CHANNELS: 256
+  HIDDEN_DIM: 512
+TRAIN:
+  BACKBONE_MULTIPLIER: 0.1
+  DROP_PATH_RATE: 0.1
+  BATCH_SIZE: 2
+  EPOCH: 180
+  GIOU_WEIGHT: 2.0
+  L1_WEIGHT: 5.0
+  QWEN_FORMAT_WEIGHT: 0.01
+  QWEN_TEACHER_WEIGHT: 0.1
+  GRAD_CLIP_NORM: 0.1
+  LR: 0.0001
+  LR_DROP_EPOCH: 144  # 4/5（0.8）
+  NUM_WORKER: 8
+  OPTIMIZER: ADAMW
+  PRINT_INTERVAL: 50
+  SCHEDULER:
+    TYPE: step
+    DECAY_RATE: 0.1
+  WEIGHT_DECAY: 0.0001
+  AMP: False
+  FIX_BN: true
+  TYPE: target_state # train target-state projector/gate and tracker head
+
+TEST:
+  EPOCH: 240
+  SEARCH_FACTOR: 4.0
+  SEARCH_SIZE: 256
+  TEMPLATE_FACTOR: 2.0
+  TEMPLATE_SIZE: 128
+  WINDOW: true
+  NUM_TEMPLATES: 2
+

install.sh

@@ -0,0 +1,108 @@
+echo "****************** Installing pytorch ******************"
+pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
+#conda install -y pytorch=1.11 torchvision torchaudio cudatoolkit=11.3 -c pytorch
+
+echo ""
+echo ""
+echo "****************** Installing yaml ******************"
+pip install PyYAML
+
+echo ""
+echo ""
+echo "****************** Installing easydict ******************"
+pip install easydict
+
+echo ""
+echo ""
+echo "****************** Installing cython ******************"
+pip install cython
+
+echo ""
+echo ""
+echo "****************** Installing opencv-python ******************"
+pip install opencv-python
+
+echo ""
+echo ""
+echo "****************** Installing pandas ******************"
+pip install pandas
+
+echo ""
+echo ""
+echo "****************** Installing tqdm ******************"
+conda install -y tqdm
+
+echo ""
+echo ""
+echo "****************** Installing coco toolkit ******************"
+pip install pycocotools
+
+echo ""
+echo ""
+echo "****************** Installing jpeg4py python wrapper ******************"
+pip install jpeg4py
+
+echo ""
+echo ""
+echo "****************** Installing tensorboard ******************"
+pip install tb-nightly
+
+echo ""
+echo ""
+echo "****************** Installing tikzplotlib ******************"
+pip install tikzplotlib
+
+echo ""
+echo ""
+echo "****************** Installing thop tool for FLOPs and Params computing ******************"
+pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git
+
+echo ""
+echo ""
+echo "****************** Installing colorama ******************"
+pip install colorama
+
+echo ""
+echo ""
+echo "****************** Installing lmdb ******************"
+pip install lmdb
+
+echo ""
+echo ""
+echo "****************** Installing scipy ******************"
+pip install scipy
+
+echo ""
+echo ""
+echo "****************** Installing visdom ******************"
+pip install visdom
+
+echo ""
+echo ""
+echo "****************** Installing vot-toolkit python ******************"
+pip install git+https://github.com/votchallenge/vot-toolkit-python
+
+echo ""
+echo ""
+echo "****************** Installing timm ******************"
+pip install timm==0.5.4
+
+echo ""
+echo ""
+echo "****************** Installing yacs ******************"
+pip install yacs
+
+echo ""
+echo ""
+echo "****************** Installing pytorch-pretrained-bert ******************"
+pip install pytorch-pretrained-bert==0.6.2
+
+echo ""
+echo ""
+echo "****************** Installing scikit-image ******************"
+pip install scikit-image
+
+echo ""
+echo ""
+
+echo "****************** Installation complete! ******************"

lib/config/atctrack/config.py

@@ -0,0 +1,184 @@
+from easydict import EasyDict as edict
+import yaml
+
+cfg = edict()
+
+# MODEL
+cfg.MODEL = edict()
+cfg.MODEL.HIDDEN_DIM = 256 # hidden dimension for the decoder and vocabulary
+cfg.MODEL.BINS = 4000 # number of discrete bins
+cfg.MODEL.FEATURE_TYPE = "x" # the input feature to decoder. x, xz, or token
+cfg.MODEL.INTERFACE_TYPE = 'low-rank_add'
+cfg.MODEL.INTERFACE_DIM = 8
+
+# MODEL.LANGUAGE
+cfg.MODEL.PRETRAIN_FILE = ""
+cfg.MODEL.PRETRAINED_PATH = ""
+
+# MODEL.TARGET_STATE
+cfg.MODEL.TARGET_STATE = edict()
+cfg.MODEL.TARGET_STATE.ENABLE = False
+cfg.MODEL.TARGET_STATE.MODEL_PATH = "../Qwen/Qwen3.5-4B"
+cfg.MODEL.TARGET_STATE.TOKEN = "<TARGET_STATE>"
+cfg.MODEL.TARGET_STATE.FREEZE_QWEN = True
+cfg.MODEL.TARGET_STATE.TRAIN_TOKEN_EMBEDDING = True
+cfg.MODEL.TARGET_STATE.USE_LORA = True
+cfg.MODEL.TARGET_STATE.LORA_R = 8
+cfg.MODEL.TARGET_STATE.LORA_ALPHA = 16
+cfg.MODEL.TARGET_STATE.LORA_DROPOUT = 0.05
+cfg.MODEL.TARGET_STATE.LORA_TARGET_MODULES = [
+    "in_proj_qkv", "out_proj", "in_proj_z", "in_proj_b", "in_proj_a",
+    "q_proj", "k_proj", "v_proj", "o_proj",
+    "gate_proj", "up_proj", "down_proj",
+]
+cfg.MODEL.TARGET_STATE.TEACHER_ENABLE = False
+cfg.MODEL.TARGET_STATE.TEACHER_MODEL = "qwen3.5"
+cfg.MODEL.TARGET_STATE.TEACHER_BASE_URL = "http://127.0.0.1:8001/v1"
+cfg.MODEL.TARGET_STATE.TEACHER_API_KEY = "sk-no-key-required"
+
+
+# MODEL.BACKBONE
+cfg.MODEL.BACKBONE = edict()
+cfg.MODEL.BACKBONE.TYPE = "vit_base_patch16_224"
+cfg.MODEL.BACKBONE.STRIDE = 16
+cfg.MODEL.BACKBONE.MID_PE = False
+cfg.MODEL.BACKBONE.SEP_SEG = False
+cfg.MODEL.BACKBONE.CAT_MODE = 'direct'
+cfg.MODEL.BACKBONE.MERGE_LAYER = 0
+cfg.MODEL.BACKBONE.ADD_CLS_TOKEN = False
+cfg.MODEL.BACKBONE.CLS_TOKEN_USE_MODE = 'ignore'
+
+# MODEL.HEAD
+cfg.MODEL.HEAD = edict()
+cfg.MODEL.HEAD.TYPE = "CENTER"
+cfg.MODEL.HEAD.NUM_CHANNELS = 512
+
+
+# TRAIN
+cfg.TRAIN = edict()
+cfg.TRAIN.LR = 0.0001
+cfg.TRAIN.WEIGHT_DECAY = 0.0001
+cfg.TRAIN.AMP = False
+cfg.TRAIN.EPOCH = 500
+cfg.TRAIN.LR_DROP_EPOCH = 400
+cfg.TRAIN.BATCH_SIZE = 8
+cfg.TRAIN.NUM_WORKER = 8
+cfg.TRAIN.OPTIMIZER = "ADAMW"
+cfg.TRAIN.ENCODER_MULTIPLIER = 0.1  # encoder's LR = this factor * LR
+cfg.TRAIN.FREEZE_ENCODER = False # for freezing the parameters of encoder
+cfg.TRAIN.ENCODER_OPEN = [] # only for debug, open some layers of encoder when FREEZE_ENCODER is True
+cfg.TRAIN.CE_WEIGHT = 1.0 # weight for cross-entropy loss
+cfg.TRAIN.GIOU_WEIGHT = 2.0
+cfg.TRAIN.L1_WEIGHT = 5.0
+cfg.TRAIN.QWEN_FORMAT_WEIGHT = 0.01
+cfg.TRAIN.QWEN_TEACHER_WEIGHT = 0.1
+
+cfg.TRAIN.PRINT_INTERVAL = 50 # interval to print the training log
+cfg.TRAIN.GRAD_CLIP_NORM = 0.1
+cfg.TRAIN.FIX_BN = False
+cfg.TRAIN.BACKBONE_MULTIPLIER = 0.1
+cfg.TRAIN.DROP_PATH_RATE = 0.1
+
+# TRAIN.SCHEDULER
+cfg.TRAIN.SCHEDULER = edict()
+cfg.TRAIN.SCHEDULER.TYPE = "step"
+cfg.TRAIN.SCHEDULER.DECAY_RATE = 0.1
+cfg.TRAIN.TYPE = "peft" # peft or fft
+cfg.TRAIN.PRETRAINED_PATH = None
+
+# DATA
+cfg.DATA = edict()
+cfg.DATA.MEAN = [0.485, 0.456, 0.406]
+cfg.DATA.STD = [0.229, 0.224, 0.225]
+cfg.DATA.MAX_SAMPLE_INTERVAL = 200
+cfg.DATA.SAMPLER_MODE = "order"
+cfg.DATA.LOADER = "tracking"
+cfg.DATA.SEQ_FORMAT = "xywh"
+cfg.DATA.MULTI_MODAL_VISION = False # vision multi-modal
+cfg.DATA.MULTI_MODAL_LANGUAGE = True # language multi-modal
+# DATA.TRAIN
+cfg.DATA.TRAIN = edict()
+cfg.DATA.TRAIN.DATASETS_NAME = ["LASOT", "GOT10K_vottrain"]
+cfg.DATA.TRAIN.DATASETS_RATIO = [1, 1]
+cfg.DATA.TRAIN.SAMPLE_PER_EPOCH = 60000
+# DATA.SEARCH
+cfg.DATA.SEARCH = edict()
+cfg.DATA.SEARCH.NUMBER = 1  #number of search region, only support 1 for now.
+cfg.DATA.SEARCH.SIZE = 256
+cfg.DATA.SEARCH.FACTOR = 4.0
+cfg.DATA.SEARCH.CENTER_JITTER = 3.5
+cfg.DATA.SEARCH.SCALE_JITTER = 0.5
+# DATA.TEMPLATE
+cfg.DATA.TEMPLATE = edict()
+cfg.DATA.TEMPLATE.NUMBER = 1
+cfg.DATA.TEMPLATE.SIZE = 256
+cfg.DATA.TEMPLATE.FACTOR = 4.0
+cfg.DATA.TEMPLATE.CENTER_JITTER = 0
+cfg.DATA.TEMPLATE.SCALE_JITTER = 0
+
+# TEST
+cfg.TEST = edict()
+cfg.TEST.TEMPLATE_FACTOR = 4.0
+cfg.TEST.TEMPLATE_SIZE = 128
+cfg.TEST.SEARCH_FACTOR = 4.0
+cfg.TEST.SEARCH_SIZE = 256
+cfg.TEST.EPOCH = 500
+cfg.TEST.WINDOW = False # window penalty
+cfg.TEST.NUM_TEMPLATES = 1
+
+cfg.TEST.UPDATE_INTERVALS = edict()
+cfg.TEST.UPDATE_INTERVALS.DEFAULT = 25
+
+cfg.TEST.UPDATE_THRESHOLD = edict()
+cfg.TEST.UPDATE_THRESHOLD.DEFAULT = 0.475
+
+cfg.TEST.MULTI_MODAL_VISION = edict()
+cfg.TEST.MULTI_MODAL_VISION.DEFAULT = False
+cfg.TEST.MULTI_MODAL_VISION.OTB99_LANG = False
+cfg.TEST.MULTI_MODAL_VISION.TNL2K = False
+cfg.TEST.MULTI_MODAL_VISION.LASOT_LANG = False
+
+
+
+
+
+def _edict2dict(dest_dict, src_edict):
+    if isinstance(dest_dict, dict) and isinstance(src_edict, dict):
+        for k, v in src_edict.items():
+            if not isinstance(v, edict):
+                dest_dict[k] = v
+            else:
+                dest_dict[k] = {}
+                _edict2dict(dest_dict[k], v)
+    else:
+        return
+
+
+def gen_config(config_file):
+    cfg_dict = {}
+    _edict2dict(cfg_dict, cfg)
+    with open(config_file, 'w') as f:
+        yaml.dump(cfg_dict, f, default_flow_style=False)
+
+
+def _update_config(base_cfg, exp_cfg):
+    if isinstance(base_cfg, dict) and isinstance(exp_cfg, edict):
+        for k, v in exp_cfg.items():
+            if k in base_cfg:
+                if not isinstance(v, dict):
+                    base_cfg[k] = v
+                else:
+                    _update_config(base_cfg[k], v)
+            else:
+                raise ValueError("{} not exist in config.py".format(k))
+    else:
+        return
+
+
+def update_config_from_file(filename):
+    exp_config = None
+    with open(filename) as f:
+        exp_config = edict(yaml.safe_load(f))
+        _update_config(cfg, exp_config)
+
+

lib/models/aqatrack/__init__.py

@@ -0,0 +1 @@
+from .aqatrack import build_aqatrack

lib/models/aqatrack/aqatrack.py

@@ -0,0 +1,196 @@
+"""
+Basic AQATrack model.
+"""
+import math
+import os
+from typing import List
+
+import torch
+from torch import nn
+from torch.nn.modules.transformer import _get_clones
+
+from lib.models.layers.head import build_box_head
+from lib.models.aqatrack.hivit import hivit_small, hivit_base
+from lib.utils.box_ops import box_xyxy_to_cxcywh
+
+from lib.models.layers.transformer_dec import build_transformer_dec
+from lib.models.layers.position_encoding import build_position_encoding
+from lib.utils.misc import NestedTensor
+
+
+import queue
+
+class AQATrack(nn.Module):
+    """ This is the base class for AQATrack """
+
+    def __init__(self, transformer, box_head, transformer_dec, position_encoding, aux_loss=False, head_type="CORNER"):
+        """ Initializes the model.
+        Parameters:
+            transformer: torch module of the transformer architecture.
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+        """
+        super().__init__()
+        self.backbone = transformer
+        self.box_head = box_head
+
+        self.aux_loss = aux_loss
+        self.head_type = head_type
+        if head_type == "CORNER" or head_type == "CENTER":
+            self.feat_sz_s = int(box_head.feat_sz)
+            self.feat_len_s = int(box_head.feat_sz ** 2)
+
+        if self.aux_loss:
+            self.box_head = _get_clones(self.box_head, 6)
+        
+        self.transformer_dec = transformer_dec
+        self.position_encoding = position_encoding 
+        self.query_embed=nn.Embedding(num_embeddings=1, embedding_dim=512)
+
+
+    def forward(self, template: torch.Tensor,
+                search: torch.Tensor,
+                return_last_attn=False,
+                training=True, #True
+                tgt_pre = None,
+                ):
+        b0, num_search = template[0].shape[0], len(search)
+        if training:
+            search = torch.cat(search, dim=0)
+            template = template[0].repeat(num_search,1,1,1)
+
+        x, aux_dict = self.backbone(z=template, x=search,
+                                    return_last_attn=return_last_attn, ) #x=[B,N,C]
+        
+        b,n,c = x.shape
+        input_dec = x
+        batches = [[] for _ in range(b0)]
+        for i, input in enumerate(input_dec):
+            batches[i % b0].append(input.unsqueeze(0))
+        x_decs = []
+        query_embed = self.query_embed.weight
+        assert len(query_embed.size()) in [2, 3]
+        if len(query_embed.size()) == 2:
+            query_embeding = query_embed.unsqueeze(1)
+        for i,batch in enumerate(batches):
+            if len(batch) ==0:
+                continue
+            tgt_all = [torch.zeros_like(query_embeding) for _ in range(num_search)]
+
+            for j, input in enumerate(batch):
+                pos_embed = self.position_encoding(1)
+                tgt_q = tgt_all[j]
+                tgt_kv = torch.cat(tgt_all[:j+1], dim=0)
+                if not training and len(tgt_pre) != 0:
+                    tgt_kv = torch.cat(tgt_pre, dim=0)
+                tgt = [tgt_q, tgt_kv]
+                tgt_out = self.transformer_dec(input.transpose(0, 1), tgt, self.feat_len_s, pos_embed, query_embeding)
+                x_decs.append(tgt_out[0])
+                tgt_all[j] = tgt_out[0]#
+            if not training:
+                if len(tgt_pre) < 3: #num_search-1 
+                    tgt_pre.append(tgt_out[0])
+                else:
+                    tgt_pre.pop(0)
+                    tgt_pre.append(tgt_out[0])
+            
+        batch0 =[]
+        if not training:
+            batch0.append(x_decs[0])
+        else:
+            batch0 = [x_decs[i + j*num_search] for j in range(b0) for i in range(num_search)]
+        
+        x_dec = torch.cat(batch0, dim = 1)
+
+        # Forward head
+        feat_last = x
+        if isinstance(x, list):
+            feat_last = x[-1]
+
+        out = self.forward_head(feat_last, x_dec, None) # STM and head
+
+        out.update(aux_dict)
+        out['tgt'] = tgt_pre 
+        #print(len(out['tgt']), 'out[tgt]')
+        return out
+
+    def forward_head(self, cat_feature, out_dec=None, gt_score_map=None):
+        """
+        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
+        """
+        # STM
+        enc_opt = cat_feature[:, -self.feat_len_s:] 
+        dec_opt = out_dec.transpose(0,1).transpose(1,2) 
+        att = torch.matmul(enc_opt, dec_opt)
+        opt = (enc_opt.unsqueeze(-1) * att.unsqueeze(-2)).permute((0, 3, 2, 1)).contiguous()
+        bs, Nq, C, HW = opt.size()
+        opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s)
+
+        #Head
+        if self.head_type == "CORNER":
+            # run the corner head
+            pred_box, score_map = self.box_head(opt_feat, True)
+            outputs_coord = box_xyxy_to_cxcywh(pred_box)
+            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
+            out = {'pred_boxes': outputs_coord_new,
+                   'score_map': score_map,
+                   }
+            return out
+
+        elif self.head_type == "CENTER":
+            # run the center head
+            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat, gt_score_map)
+            # outputs_coord = box_xyxy_to_cxcywh(bbox)
+            outputs_coord = bbox
+            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
+            out = {'pred_boxes': outputs_coord_new,
+                   'score_map': score_map_ctr,
+                   'size_map': size_map,
+                   'offset_map': offset_map}
+            return out
+        else:
+            raise NotImplementedError
+
+
+def build_aqatrack(cfg, training=True):
+    current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
+    pretrained_path = os.path.join(current_dir, '../../../pretrained_models')
+    if cfg.MODEL.PRETRAIN_FILE and ('AQATrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
+        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
+    else:
+        pretrained = ''
+
+    if cfg.MODEL.BACKBONE.TYPE == 'hivit_small':
+        backbone = hivit_small(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+
+    elif cfg.MODEL.BACKBONE.TYPE == 'hivit_base':
+        backbone = hivit_base(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+
+    else:
+        raise NotImplementedError
+
+    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)
+
+    
+    transformer_dec = build_transformer_dec(cfg, hidden_dim)
+    position_encoding = build_position_encoding(cfg, sz = 1)
+
+    box_head = build_box_head(cfg, hidden_dim)
+    model = AQATrack(
+        backbone,
+        box_head,
+        transformer_dec,
+        position_encoding,        
+        aux_loss=False,
+        head_type=cfg.MODEL.HEAD.TYPE,
+    )
+
+    if 'AQATrack' in cfg.MODEL.PRETRAIN_FILE and training:
+        checkpoint = torch.load(cfg.MODEL.PRETRAIN_FILE, map_location="cpu")
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint["net"], strict=False)
+        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)
+
+    return model

lib/models/aqatrack/base_backbone.py

@@ -0,0 +1,142 @@
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.vision_transformer import resize_pos_embed
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+from lib.models.layers.patch_embed import PatchEmbed
+from lib.models.aqatrack.utils import combine_tokens, recover_tokens
+
+
+class BaseBackbone(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.pos_embed = None
+        self.img_size = [224, 224]
+        self.patch_size = 16
+        self.embed_dim = 384
+
+        self.cat_mode = 'direct'
+
+        self.pos_embed_z = None
+        self.pos_embed_x = None
+
+        self.template_segment_pos_embed = None
+        self.search_segment_pos_embed = None
+
+        self.return_inter = False
+        self.return_stage = [2, 5, 8, 11]
+
+        self.add_cls_token = True
+        self.add_sep_seg = False
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, 512))
+
+    def finetune_track(self, cfg,dim, patch_start_index=1):
+
+        search_size = to_2tuple(cfg.DATA.SEARCH.SIZE)
+        template_size = to_2tuple(cfg.DATA.TEMPLATE.SIZE)
+        new_patch_size = cfg.MODEL.BACKBONE.STRIDE
+
+        self.cat_mode = cfg.MODEL.BACKBONE.CAT_MODE
+        self.return_inter = False
+
+        patch_pos_embed_all = self.pos_embed
+        patch_pos_embed = patch_pos_embed_all[:,1:,:]
+        patch_pos_embed = patch_pos_embed.transpose(1, 2)
+        B, E, Q = patch_pos_embed.shape
+        P_H, P_W = self.img_size // self.patch_size, self.img_size // self.patch_size
+        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
+
+        # for search region
+        H, W = search_size
+        new_P_H, new_P_W = H // new_patch_size, W // new_patch_size
+        search_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H, new_P_W), mode='bicubic',
+                                                           align_corners=False)
+        search_patch_pos_embed = search_patch_pos_embed.flatten(2).transpose(1, 2)
+
+        # for template region
+        H, W = template_size
+        new_P_H, new_P_W = H // new_patch_size, W // new_patch_size
+        template_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H, new_P_W), mode='bicubic',
+                                                             align_corners=False)
+        template_patch_pos_embed = template_patch_pos_embed.flatten(2).transpose(1, 2)
+
+        self.pos_embed_z = nn.Parameter(torch.cat([template_patch_pos_embed,template_patch_pos_embed],dim=1))
+        self.pos_embed_x = nn.Parameter(search_patch_pos_embed)
+
+        cls_token = patch_pos_embed_all[:,0,:].unsqueeze(1)
+        self.cls_token = nn.Parameter(cls_token)
+
+        # token_type
+        self.token_type_search = nn.Parameter(torch.zeros(1, 1, dim))
+        self.token_type_template_fg = nn.Parameter(torch.zeros(1, 1, dim))
+        self.token_type_template_bg = nn.Parameter(torch.zeros(1, 1, dim))
+
+
+        if self.return_inter:
+            for i_layer in self.fpn_stage:
+                if i_layer != 11:
+                    norm_layer = partial(nn.LayerNorm, eps=1e-6)
+                    layer = norm_layer(self.embed_dim)
+                    layer_name = f'norm{i_layer}'
+                    self.add_module(layer_name, layer)
+
+    def forward_features(self, z, x, mask=None):
+        B = x.shape[0]
+    
+        z = self.patch_embed(z)
+        x = self.patch_embed(x)
+    
+        for blk in self.blocks[:-self.num_main_blocks]:
+            x = blk(x)
+            z = blk(z)
+
+        x = x[..., 0, 0, :]
+        z = z[..., 0, 0, :]
+
+
+        z += self.pos_embed_z
+        x += self.pos_embed_x
+
+        lens_z = self.pos_embed_z.shape[1]
+        lens_x = self.pos_embed_x.shape[1]
+        
+        x = combine_tokens(z, x, mode=self.cat_mode)
+        #x = combine_tokens(x, z,  mode=self.cat_mode)
+
+
+        if self.add_cls_token:
+            cls_tokens = self.cls_token.expand(B, -1, -1)
+            # cls_tokens = cls_tokens + self.cls_pos_embed
+            x = torch.cat([cls_tokens, x], dim=1)
+
+        x = self.pos_drop(x)
+
+        for blk in self.blocks[-self.num_main_blocks:]:
+            x = blk(x)
+
+        x = recover_tokens(x, lens_z, lens_x, mode=self.cat_mode)
+
+        aux_dict = {"attn": None}
+        x = self.norm_(x)
+
+        return x, aux_dict
+
+    def forward(self, z, x, **kwargs):
+        """
+        Joint feature extraction and relation modeling for the basic HiViT backbone.
+        Args:
+            z (torch.Tensor): template feature, [B, C, H_z, W_z]
+            x (torch.Tensor): search region feature, [B, C, H_x, W_x]
+
+        Returns:
+            x (torch.Tensor): merged template and search region feature, [B, L_z+L_x, C]
+            attn : None
+        """
+        x, aux_dict = self.forward_features(z, x,)
+
+        return x, aux_dict

lib/models/aqatrack/fast_itpn.py

@@ -0,0 +1,1173 @@
+# --------------------------------------------------------
+# Fast-iTPN: Integrally Pre-Trained Transformer Pyramid Network with Token Migration
+# Github source: https://github.com/sunsmarterjie/iTPN/tree/main/fast_itpn
+# Copyright (c) 2023 University of Chinese Academy of Sciences
+# Licensed under The MIT License [see LICENSE for details]
+# By Yunjie Tian
+# Based on EVA02, timm and deit code bases
+# https://github.com/baaivision/EVA/tree/master/EVA-02
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# https://github.com/facebookresearch/deit/
+# --------------------------------------------------------'
+from functools import partial
+
+import math
+import torch
+import torch.nn as nn
+from timm.models.registry import register_model
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import to_2tuple, drop_path, trunc_normal_
+
+from torch import Tensor, Size
+from typing import Union, List
+from lib.models.aqatrack.base_backbone import BaseBackbone
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
+        **kwargs
+    }
+
+
+_shape_t = Union[int, List[int], Size]
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.,
+                 norm_layer=nn.LayerNorm, subln=False
+                 ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.ffn_ln(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class ConvMlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.,
+                 norm_layer=nn.LayerNorm, subln=False
+                 ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else None
+
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        if self.ffn_ln is not None:
+            x = x.permute(0, 2, 3, 1)
+            x = self.ffn_ln(x)
+            x = x.permute(0, 3, 1, 2)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
+                 norm_layer=nn.LayerNorm, subln=False
+                 ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Linear(hidden_features, out_features)
+
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+
+class ConvSwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
+                 norm_layer=nn.LayerNorm, subln=False
+                 ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.w2 = nn.Conv2d(in_features, hidden_features, 1)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Conv2d(hidden_features, out_features, 1)
+
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x1 = self.w1(x).flatten(2).transpose(1, 2)
+        x2 = self.w2(x).flatten(2).transpose(1, 2)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden).transpose(1, 2).view(B, C, H, W)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., window_size=None,
+            attn_head_dim=None, use_decoupled_rel_pos_bias=False, deepnorm=False, subln=False
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.deepnorm = deepnorm
+        self.subln = subln
+        if self.deepnorm or self.subln:
+            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+        else:
+            self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.rel_pos_bias = None
+        self.qk_float = True
+
+        self.window_size = None
+        self.relative_position_bias_table = None
+
+        if window_size:
+            if use_decoupled_rel_pos_bias:
+                self.rel_pos_bias = DecoupledRelativePositionBias(window_size=window_size, num_heads=num_heads)
+            else:
+                self.window_size = window_size
+                self.num_relative_distance = (2 * window_size[0] - 1) * (
+                        2 * window_size[1] - 1) + 3  # (2*14-1) * (2*14-1) + 3
+                self.relative_position_bias_table = nn.Parameter(
+                    torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+                # cls to token & token 2 cls & cls to cls
+
+                # get pair-wise relative position index for each token inside the window
+                coords_h = torch.arange(window_size[0])
+                coords_w = torch.arange(window_size[1])
+                coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+                coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+                relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+                relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+                relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+                relative_coords[:, :, 1] += window_size[1] - 1
+                relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+                relative_position_index = \
+                    torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+                relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+                relative_position_index[0, 0:] = self.num_relative_distance - 3
+                relative_position_index[0:, 0] = self.num_relative_distance - 2
+                relative_position_index[0, 0] = self.num_relative_distance - 1
+
+                self.register_buffer("relative_position_index", relative_position_index)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None,require_attn = False):
+        B, N, C = x.shape
+
+        if self.deepnorm or self.subln:
+            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+        else:
+            qkv_bias = None
+            if self.q_bias is not None:
+                qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+            qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)  # 3, B, num_heads, N, C
+            q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q * self.scale
+        if self.qk_float:
+            attn = (q.float() @ k.float().transpose(-2, -1))
+        else:
+            attn = (q @ k.transpose(-2, -1))
+
+        if self.relative_position_bias_table is not None:
+            relative_position_bias = \
+                self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                    self.window_size[0] * self.window_size[1] + 1,
+                    self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+            attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+        if self.rel_pos_bias is not None:
+            attn = attn + self.rel_pos_bias().type_as(attn)
+
+        if rel_pos_bias is not None:
+            attn = attn + rel_pos_bias.type_as(attn)
+        if attn_mask is not None:
+            attn_mask = attn_mask.bool()
+            attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+        attn = attn.softmax(dim=-1).type_as(x)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        if require_attn:
+            return x,attn
+        else:
+            return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, norm_layer=nn.LayerNorm, window_size=None, attn_head_dim=None,
+                 use_decoupled_rel_pos_bias=False,
+                 depth=None,
+                 postnorm=False,
+                 deepnorm=False,
+                 subln=False,
+                 swiglu=False,
+                 naiveswiglu=False,
+                 ):
+        super().__init__()
+
+        with_attn = num_heads > 0
+
+        self.norm1 = norm_layer(dim) if with_attn else None
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size,
+            use_decoupled_rel_pos_bias=use_decoupled_rel_pos_bias, attn_head_dim=attn_head_dim,
+            deepnorm=deepnorm,
+            subln=subln
+        ) if with_attn else None
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        if swiglu:
+            self.mlp = xops.SwiGLU(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim
+            )  # hidden_features: 2/3
+        elif naiveswiglu:
+            self.mlp = SwiGLU(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = Mlp(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                subln=subln,
+                norm_layer=norm_layer
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),
+                                        requires_grad=True) if self.attn is not None else None
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.deepnorm = deepnorm
+        if self.deepnorm:
+            self.alpha = math.pow(2.0 * depth, 0.25)
+
+        self.postnorm = postnorm
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None,require_attn= False):
+        if self.gamma_2 is None:
+            if self.postnorm:
+                if self.attn is not None:
+                    x = x + self.drop_path(
+                        self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            elif self.deepnorm:
+                if self.attn is not None:
+                    residual = x
+                    x = self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)
+                    x = self.drop_path(x)
+                    x = residual * self.alpha + x
+                    x = self.norm1(x)
+
+                residual = x
+                x = self.mlp(x)
+                x = self.drop_path(x)
+                x = residual * self.alpha + x
+                x = self.norm2(x)
+            else:
+                if self.attn is not None:
+                    x = x + self.drop_path(
+                        self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.postnorm:
+                if self.attn is not None:
+                    x = x + self.drop_path(
+                        self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                if self.attn is not None:
+                    if require_attn:
+                        x_m,attn = self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask, require_attn=True)
+                        x = x + self.drop_path(
+                            self.gamma_1 *x_m )
+                    else:
+                        x = x + self.drop_path(
+                            self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        if require_attn:
+            return x,attn
+        else:
+            return x
+
+
+class ConvMlpBlock(nn.Module):
+
+    def __init__(self, dim, mlp_ratio=4., drop_path=0., init_values=None, norm_layer=nn.LayerNorm,
+                 depth=None,
+                 postnorm=False,
+                 deepnorm=False,
+                 subln=False,
+                 swiglu=False,
+                 naiveswiglu=False,
+                 ):
+        super().__init__()
+
+        self.attn = None
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        if swiglu:
+            self.mlp = xops.SwiGLU(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim
+            )  # hidden_features: 2/3
+        elif naiveswiglu:
+            self.mlp = ConvSwiGLU(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = ConvMlp(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                subln=subln,
+                norm_layer=norm_layer
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),
+                                        requires_grad=True) if self.attn is not None else None
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.deepnorm = deepnorm
+        if self.deepnorm:
+            self.alpha = math.pow(2.0 * depth, 0.25)
+
+        self.postnorm = postnorm
+
+    def forward(self, x):
+        if self.gamma_2 is None:
+            if self.postnorm:
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            elif self.deepnorm:
+                residual = x
+                x = self.mlp(x)
+                x = self.drop_path(x)
+                x = residual * self.alpha + x
+                x = self.norm2(x)
+            else:
+                x = x + self.drop_path(self.mlp(self.norm2(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)))
+        else:
+            if self.postnorm:
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                m = self.mlp(self.norm2(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2))
+                x = x + self.drop_path(self.gamma_2 * m)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, inner_patches=4, in_chans=3, embed_dim=128, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.inner_patches = inner_patches
+        self.patches_resolution = self.patch_shape = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        # self.in_chans = 6
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        conv_size = [size // inner_patches for size in patch_size]
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=conv_size, stride=conv_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        patches_resolution = (H // self.patch_size[0], W // self.patch_size[1])
+        num_patches = patches_resolution[0] * patches_resolution[1]
+        x = self.proj(x).view(
+            B, -1,
+            patches_resolution[0], self.inner_patches,
+            patches_resolution[1], self.inner_patches,
+        ).permute(0, 2, 4, 3, 5, 1).reshape(B, num_patches, self.inner_patches, self.inner_patches, -1)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+class ConvPatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, inner_patches=4, in_chans=3, embed_dim=128, norm_layer=None,
+                 stop_grad_conv1=False):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.stop_grad_conv1 = stop_grad_conv1
+        self.inner_patches = inner_patches
+        self.patches_resolution = self.patch_shape = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        # self.in_chans = 6
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        conv_size = [size // inner_patches for size in patch_size]
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=conv_size, stride=conv_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x, bool_masked_pos=None, mask_token=None):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.stop_grad_conv1:
+            x = x.detach() * 0.9 + x * 0.1
+
+        if bool_masked_pos is not None:
+            x = torch.nn.functional.unfold(x, kernel_size=4, stride=4, padding=0).transpose(1, 2)
+
+            seq_len = x.shape[1]
+            mask_token = mask_token.expand(B, seq_len, -1)
+            w = bool_masked_pos.unsqueeze(-1).type_as(mask_token)
+            x = x * (1 - w) + mask_token * w
+
+            x = torch.nn.functional.fold(x.transpose(1, 2), output_size=(H // 4, W // 4), kernel_size=4, padding=0,
+                                         stride=4)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+class PatchMerge(nn.Module):
+    def __init__(self, dim, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(dim * 4)
+        self.reduction = nn.Linear(dim * 4, dim * 2, bias=False)
+        self.mlp = None
+
+    def forward(self, x):
+        x0 = x[..., 0::2, 0::2, :]
+        x1 = x[..., 1::2, 0::2, :]
+        x2 = x[..., 0::2, 1::2, :]
+        x3 = x[..., 1::2, 1::2, :]
+
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+        x = self.norm(x)
+        x = self.reduction(x)
+        return x
+
+
+class ConvPatchMerge(nn.Module):
+    def __init__(self, dim, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(dim)
+        self.reduction = nn.Conv2d(dim, dim * 2, kernel_size=2, stride=2, padding=0)
+        self.mlp = None
+
+    def forward(self, x):
+        x = self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        x = self.reduction(x)
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = \
+            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def forward(self):
+        relative_position_bias = \
+            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                self.window_size[0] * self.window_size[1] + 1,
+                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+
+def _mask_1d_rel_pos_index(seq_len):
+    index = torch.arange(seq_len)
+    return index.view(1, seq_len) - index.view(seq_len, 1) + seq_len - 1
+
+
+def _add_cls_to_index_matrix(index, num_tokens, offset):
+    index = index.contiguous().view(num_tokens, num_tokens)
+    new_index = torch.zeros(size=(num_tokens + 1, num_tokens + 1), dtype=index.dtype)
+    new_index[1:, 1:] = index
+    new_index[0, 0:] = offset
+    new_index[0:, 0] = offset + 1
+    new_index[0, 0] = offset + 2
+    return new_index
+
+
+class DecoupledRelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] + 2, 2 * window_size[1] + 2)
+
+        num_tokens = window_size[0] * window_size[1]
+
+        self.relative_position_bias_for_high = nn.Parameter(torch.zeros(self.num_relative_distance[0], num_heads))
+        self.relative_position_bias_for_width = nn.Parameter(torch.zeros(self.num_relative_distance[1], num_heads))
+        # cls to token & token 2 cls & cls to cls
+
+        h_index = _mask_1d_rel_pos_index(window_size[0]).view(
+            window_size[0], 1, window_size[0], 1).expand(-1, window_size[1], -1, window_size[1])
+        h_index = _add_cls_to_index_matrix(h_index, num_tokens, 2 * window_size[0] - 1)
+        self.register_buffer("relative_position_high_index", h_index)
+
+        w_index = _mask_1d_rel_pos_index(window_size[1]).view(
+            1, window_size[1], 1, window_size[1]).expand(window_size[0], -1, window_size[0], -1)
+        w_index = _add_cls_to_index_matrix(w_index, num_tokens, 2 * window_size[1] - 1)
+
+        self.register_buffer("relative_position_width_index", w_index)
+
+    def forward(self):
+        relative_position_bias = \
+            F.embedding(input=self.relative_position_high_index, weight=self.relative_position_bias_for_high) + \
+            F.embedding(input=self.relative_position_width_index, weight=self.relative_position_bias_for_width)
+        return relative_position_bias.permute(2, 0, 1).contiguous()
+
+
+class Fast_iTPN(BaseBackbone):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=512, depth_stage1=3, depth_stage2=3, depth=24,
+                 num_heads=8, bridge_mlp_ratio=3., mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0.0, init_values=0.1, attn_head_dim=None, norm_layer=nn.LayerNorm,
+                 patch_norm=False, num_classes=1000, use_mean_pooling=False,
+                 init_scale=0.01,
+                 cls_token=False,
+                 grad_ckpt=False,
+                 stop_grad_conv1=False,
+                 use_abs_pos_emb=True,
+                 use_rel_pos_bias=False,
+                 use_shared_rel_pos_bias=False,
+                 use_shared_decoupled_rel_pos_bias=False,
+                 convmlp=False,
+
+                 postnorm=False,
+                 deepnorm=False,
+                 subln=False,
+                 swiglu=False,
+                 naiveswiglu=False,
+                 **kwargs):
+        super().__init__()
+        self.img_size = img_size
+        self.mlp_ratio = mlp_ratio
+        self.grad_ckpt = grad_ckpt
+        self.num_main_blocks = depth
+        self.depth_stage1 = depth_stage1
+        self.depth_stage2 = depth_stage2
+        self.depth = depth
+        self.patch_size = patch_size
+        self.num_features = self.embed_dim = embed_dim
+        self.convmlp = convmlp
+        self.stop_grad_conv1 = stop_grad_conv1
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.use_shared_rel_pos_bias = use_shared_rel_pos_bias
+        self.use_shared_decoupled_rel_pos_bias = use_shared_decoupled_rel_pos_bias
+        self.use_decoupled_rel_pos_bias = False
+
+        mlvl_dims = {'4': embed_dim // 4, '8': embed_dim // 2, '16': embed_dim}
+        # split image into non-overlapping patches
+        # in_chans = 6
+        if convmlp:
+            self.patch_embed = ConvPatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=mlvl_dims['4'],
+                stop_grad_conv1=stop_grad_conv1,
+                norm_layer=norm_layer if patch_norm else None)
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=mlvl_dims['4'],
+                norm_layer=norm_layer if patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+
+        if cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        else:
+            self.cls_token = None
+        if use_abs_pos_emb:
+            if cls_token:
+                self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+            else:
+                self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        if use_shared_rel_pos_bias:
+            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+        else:
+            self.rel_pos_bias = None
+
+        if use_shared_decoupled_rel_pos_bias:
+            assert self.rel_pos_bias is None
+            self.rel_pos_bias = DecoupledRelativePositionBias(window_size=self.patch_embed.patch_shape,
+                                                              num_heads=num_heads)
+
+        self.subln = subln
+        self.swiglu = swiglu
+        self.naiveswiglu = naiveswiglu
+
+        self.build_blocks(
+            depths=[depth_stage1, depth_stage2, depth],
+            dims=mlvl_dims,
+            num_heads=num_heads,
+            bridge_mlp_ratio=bridge_mlp_ratio,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+            drop=drop_rate,
+            attn_drop=attn_drop_rate,
+            drop_path_rate=drop_path_rate,
+            norm_layer=norm_layer,
+            init_values=init_values,
+            attn_head_dim=attn_head_dim,
+            postnorm=postnorm,
+            deepnorm=deepnorm,
+            subln=subln,
+            swiglu=swiglu,
+            naiveswiglu=naiveswiglu,
+            convmlp=convmlp,
+        )
+
+        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=.02)
+        if self.cls_token is not None:
+            trunc_normal_(self.cls_token, std=.02)
+
+        self.apply(self._init_weights)
+
+
+    def build_blocks(self,
+                     depths=[3, 3, 24],
+                     dims={'4': 128 // 4, '8': 256, '16': 512},
+                     num_heads=8,
+                     bridge_mlp_ratio=3.,
+                     mlp_ratio=4.0,
+                     qkv_bias=True,
+                     qk_scale=None,
+                     window_size=None,
+                     drop=0.,
+                     attn_drop=0.,
+                     drop_path_rate=0.,
+                     norm_layer=nn.LayerNorm,
+                     init_values=0.,
+                     attn_head_dim=None,
+                     postnorm=False,
+                     deepnorm=False,
+                     subln=False,
+                     swiglu=False,
+                     naiveswiglu=False,
+                     convmlp=False,
+                     ):
+        dpr = iter(x.item() for x in torch.linspace(0, drop_path_rate, depths[0] + depths[1] + depths[2]))
+
+        self.blocks = nn.ModuleList()
+
+        if convmlp:
+            self.blocks.extend([
+                ConvMlpBlock(
+                    dim=dims['4'],
+                    mlp_ratio=bridge_mlp_ratio,
+                    drop_path=next(dpr),
+                    norm_layer=norm_layer,
+                    init_values=0.,
+                    depth=depths[-1],
+                    postnorm=postnorm,
+                    deepnorm=deepnorm,
+                    subln=subln,
+                    swiglu=False,
+                    naiveswiglu=False,
+                ) for _ in range(depths[0])
+            ])
+            self.blocks.append(ConvPatchMerge(dims['4'], norm_layer))
+            self.blocks.extend([
+                ConvMlpBlock(
+                    dim=dims['8'],
+                    mlp_ratio=bridge_mlp_ratio,
+                    drop_path=next(dpr),
+                    norm_layer=norm_layer,
+                    init_values=0.,
+                    depth=depths[-1],
+                    postnorm=postnorm,
+                    deepnorm=deepnorm,
+                    subln=subln,
+                    swiglu=False,
+                    naiveswiglu=False,
+                ) for _ in range(depths[1])
+            ])
+            self.blocks.append(ConvPatchMerge(dims['8'], norm_layer))
+        else:
+            self.blocks.extend([
+                Block(
+                    dim=dims['4'],
+                    num_heads=0,
+                    mlp_ratio=bridge_mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=next(dpr),
+                    norm_layer=norm_layer,
+                    init_values=init_values,
+                    window_size=window_size,
+                    depth=depths[-1],
+                    postnorm=postnorm,
+                    deepnorm=deepnorm,
+                    subln=subln,
+                    swiglu=swiglu,
+                    naiveswiglu=naiveswiglu,
+                ) for _ in range(depths[0])
+            ])
+            self.blocks.append(PatchMerge(dims['4'], norm_layer))
+            self.blocks.extend([
+                Block(
+                    dim=dims['8'],
+                    num_heads=0,
+                    mlp_ratio=bridge_mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=next(dpr),
+                    norm_layer=norm_layer,
+                    init_values=init_values,
+                    window_size=window_size,
+                    depth=depths[-1],
+                    postnorm=postnorm,
+                    deepnorm=deepnorm,
+                    subln=subln,
+                    swiglu=swiglu,
+                    naiveswiglu=naiveswiglu,
+                ) for _ in range(depths[1])
+            ])
+            self.blocks.append(PatchMerge(dims['8'], norm_layer))
+
+        ######### stage 3 ########
+        self.blocks.extend([
+            Block(
+                dim=dims['16'],
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=next(dpr),
+                norm_layer=norm_layer,
+                init_values=init_values,
+                window_size=window_size,
+                attn_head_dim=attn_head_dim,
+                depth=depths[-1],
+                postnorm=postnorm,
+                deepnorm=deepnorm,
+                subln=subln,
+                swiglu=swiglu,
+                naiveswiglu=naiveswiglu,
+            ) for _ in range(depths[2])
+        ])
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        if self.cls_token is not None:
+            return {'pos_embed', 'cls_token'}
+        return {'pos_embed'}
+
+    def get_classifer(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+    def forward_features_pe(self, z,x,soft_token_template_mask):
+        B, C, H, W = x.shape
+        z, z1 = z[0], z[1]
+        z = self.patch_embed(z)
+        z1 = self.patch_embed(z1)
+        x = self.patch_embed(x)
+
+        for blk in self.blocks[:-self.num_main_blocks]:
+            x = blk(x)
+            z = blk(z)
+            z1 = blk(z1)
+
+        x = x.flatten(2).transpose(1, 2)
+        z = z.flatten(2).transpose(1, 2)
+        z1 = z1.flatten(2).transpose(1, 2)
+
+        token_type_search = self.token_type_search.expand(B, x.shape[1], -1)
+        token_type_template_bg = self.token_type_template_bg.expand(B, z.shape[1], -1)
+        token_type_template_fg = self.token_type_template_fg.expand(B, z.shape[1], -1)
+
+        x += token_type_search
+
+        token_type_template_fg = soft_token_template_mask[0] * token_type_template_fg
+        token_type_template_bg = (1 - soft_token_template_mask[0]) * token_type_template_bg
+        token_type_template_0 = token_type_template_fg + token_type_template_bg
+        z += token_type_template_0
+
+        token_type_template_fg = soft_token_template_mask[1] * token_type_template_fg
+        token_type_template_bg = (1 - soft_token_template_mask[1]) * token_type_template_bg
+        token_type_template_1 = token_type_template_fg + token_type_template_bg
+        z1 += token_type_template_1
+
+        x = torch.cat([z, z1, x], dim=1)
+
+        return x,[token_type_search,token_type_template_0,token_type_template_1 ]
+
+    def forward_features_stage3(self, x, update_cls_token,pos_embed=None):
+        B, L, D = x.shape
+        if self.cls_token is not None:
+            if update_cls_token is None:
+                cls_tokens = self.cls_token.expand(B, -1, -1)
+            else:
+                cls_tokens = update_cls_token
+            x = torch.cat([cls_tokens, x], dim=1)
+
+        if pos_embed is None:
+            pos_embed_z = self.pos_embed_z.expand(B, -1, -1)
+            pos_embed_x = self.pos_embed_x.expand(B, -1, -1)
+            cls_token = self.cls_token.expand(B, -1, -1)
+            pos_embed = torch.cat([cls_token, pos_embed_z, pos_embed_x], dim=1)
+
+        x = x + pos_embed
+        x = self.pos_drop(x)
+
+        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+        for blk in self.blocks[-self.num_main_blocks:]:
+            x = blk(x, rel_pos_bias)
+
+        x = self.norm(x)
+        aux_dict = {"attn": None}
+        return x, aux_dict
+
+    def forward_features(self, z,x,update_cls_token,soft_token_template_mask):
+        B, C, H, W = x.shape
+        z,z1 = z[0],z[1]
+        z = self.patch_embed(z)
+        z1 = self.patch_embed(z1)
+        x = self.patch_embed(x)
+
+
+        for blk in self.blocks[:-self.num_main_blocks]:
+            x =  blk(x)
+            z = blk(z)
+            z1 = blk(z1)
+
+        x = x.flatten(2).transpose(1, 2)
+        z = z.flatten(2).transpose(1, 2)
+        z1 = z1.flatten(2).transpose(1, 2)
+
+        token_type_search = self.token_type_search.expand(B, x.shape[1], -1)
+        token_type_template_bg = self.token_type_template_bg.expand(B, z.shape[1], -1)
+        token_type_template_fg = self.token_type_template_fg.expand(B, z.shape[1], -1)
+
+        x += token_type_search
+
+        token_type_template_fg = soft_token_template_mask[0] * token_type_template_fg
+        token_type_template_bg = (1 - soft_token_template_mask[0]) * token_type_template_bg
+        z += token_type_template_fg + token_type_template_bg
+
+        token_type_template_fg = soft_token_template_mask[1] * token_type_template_fg
+        token_type_template_bg = (1 - soft_token_template_mask[1]) * token_type_template_bg
+        z1 += token_type_template_fg + token_type_template_bg
+
+        # soft_token_debug_1, soft_token_debug_2 = soft_token_template_mask[0][0,:,0].view(7,7),soft_token_template_mask[1][0,:,0].view(7,7)
+
+        if self.cls_token is not None:
+            if update_cls_token is None:
+                cls_tokens = self.cls_token.expand(B, -1, -1)
+            else:
+                cls_tokens = update_cls_token
+            x = torch.cat([cls_tokens, z,z1,x], dim=1)
+
+        pos_embed_z = self.pos_embed_z.expand(B, -1, -1)
+        pos_embed_x = self.pos_embed_x.expand(B, -1, -1)
+        cls_token = self.cls_token.expand(B, -1, -1)
+        pos_embed = torch.cat([cls_token,pos_embed_z,pos_embed_x],dim=1)
+
+
+        x = x + pos_embed
+        x = self.pos_drop(x)
+
+        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+        for blk in self.blocks[-self.num_main_blocks:]:
+            x,attn =  blk(x, rel_pos_bias,require_attn= True)
+
+        x = self.norm(x)
+        aux_dict = {"attn": attn}
+        return x, aux_dict
+
+    def forward(self, z,x,update_cls_token,soft_token_template_mask):
+        x, aux_dict = self.forward_features(z, x, update_cls_token,soft_token_template_mask)
+        return x, aux_dict
+
+
+@register_model
+def fast_itpn_tiny_1112_patch16_224(pretrained=False, **kwargs):
+    model = Fast_iTPN(
+        patch_size=16, embed_dim=384, depth_stage1=1, depth_stage2=1, depth=12, num_heads=6, bridge_mlp_ratio=3.,
+        mlp_ratio=3., qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        convmlp=True,
+        naiveswiglu=True,
+        subln=True,
+        **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        checkpoint = torch.load(
+            kwargs["init_ckpt"], map_location="cpu"
+        )
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+
+@register_model
+def fast_itpn_small_2220_patch16_224(pretrained=False, **kwargs):
+    model = Fast_iTPN(
+        patch_size=16, embed_dim=384, depth_stage1=2, depth_stage2=2, depth=20, num_heads=6, bridge_mlp_ratio=3.,
+        mlp_ratio=3., qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        convmlp=True,
+        naiveswiglu=True,
+        subln=True,
+        **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        checkpoint = torch.load(
+            kwargs["init_ckpt"], map_location="cpu"
+        )
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+
+@register_model
+def fast_itpn_base_3324_patch16_224(pretrained="", **kwargs):
+    model = Fast_iTPN(
+        patch_size=16, embed_dim=512, depth_stage1=3, depth_stage2=3, depth=24, num_heads=8, bridge_mlp_ratio=3.,
+        mlp_ratio=3., qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        convmlp=True,
+        naiveswiglu=True,
+        subln=True,
+        cls_token = True,           # for task cls
+        **kwargs)
+    model.default_cfg = _cfg()
+
+    if pretrained:
+        checkpoint = torch.load(
+            pretrained, map_location="cpu"
+        )
+        checkpoint = checkpoint["module"]
+        ###### naive method
+        # missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+
+        #### process patch_embed.proj
+        normal_weights = {}
+        # special_weights = {}
+        for k, v in checkpoint.items():
+            # if "patch_embed.proj.weight" in k:
+            #     v1 = torch.cat([v, v], dim=1)*0.5
+            #     normal_weights[k] = v1
+            if "pos_embed" in k:
+                # add cls token
+                cls_pos = torch.mean(v,dim=1).unsqueeze(dim=1)
+                v1 = torch.cat([cls_pos,v],dim=1)
+                normal_weights[k] = v1
+
+            else:
+                normal_weights[k] = v
+
+        missing_keys, unexpected_keys = model.load_state_dict(normal_weights, strict=False)
+
+
+    return model
+
+@register_model
+def fast_itpn_large_2240_patch16_256(pretrained=False, **kwargs):
+    model = Fast_iTPN(
+        patch_size=16, embed_dim=768, depth_stage1=2, depth_stage2=2, depth=40, num_heads=12, bridge_mlp_ratio=3.,
+        mlp_ratio=3., qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        convmlp=True,
+        naiveswiglu=True,
+        subln=True,
+        cls_token=True,  # for task cls
+        **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        checkpoint = torch.load(
+            pretrained, map_location="cpu"
+        )
+        checkpoint = checkpoint["module"]
+        ###### naive method
+        # missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+
+        #### process patch_embed.proj
+        normal_weights = {}
+        # special_weights = {}
+        for k, v in checkpoint.items():
+            # if "patch_embed.proj.weight" in k:
+            #     v1 = torch.cat([v, v], dim=1)*0.5
+            #     normal_weights[k] = v1
+            if "pos_embed" in k:
+                # add cls token
+                cls_pos = torch.mean(v, dim=1).unsqueeze(dim=1)
+                v1 = torch.cat([cls_pos, v], dim=1)
+                normal_weights[k] = v1
+
+            else:
+                normal_weights[k] = v
+
+        missing_keys, unexpected_keys = model.load_state_dict(normal_weights, strict=False)
+
+    return model

lib/models/aqatrack/hivit.py

@@ -0,0 +1,270 @@
+""" Vision Transformer (ViT) in PyTorch
+A PyTorch implement of Vision Transformers as described in:
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale'
+    - https://arxiv.org/abs/2010.11929
+`How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers`
+    - https://arxiv.org/abs/2106.10270
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+DeiT model defs and weights from https://github.com/facebookresearch/deit,
+paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+Hacked together by / Copyright 2021 Ross Wightman
+
+Modified by Botao Ye
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+import torch.utils.checkpoint as checkpoint
+from timm.models.vision_transformer import DropPath, Mlp, trunc_normal_
+from timm.models.layers import to_2tuple
+from lib.models.aqatrack.base_backbone import BaseBackbone
+
+class Attention(nn.Module):
+    def __init__(self, input_size, dim, num_heads, qkv_bias=True, qk_scale=None, 
+                 attn_drop=0., proj_drop=0., rpe=True):
+        super().__init__()
+        self.input_size = input_size
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, rpe_index=None, mask=None):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        attn = attn.float().clamp(min=torch.finfo(torch.float32).min, max=torch.finfo(torch.float32).max)
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class BlockWithRPE(nn.Module):
+    def __init__(self, input_size, dim, num_heads=0., mlp_ratio=4., qkv_bias=True, qk_scale=None, 
+                 drop=0., attn_drop=0., drop_path=0., rpe=True,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+
+        with_attn = num_heads > 0.
+
+        self.norm1 = norm_layer(dim) if with_attn else None
+        self.attn = Attention(
+            input_size, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
+            attn_drop=attn_drop, proj_drop=drop, rpe=rpe,
+        ) if with_attn else None
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, rpe_index=None, mask=None):
+        if self.attn is not None:
+            x = x + self.drop_path(self.attn(self.norm1(x), rpe_index, mask))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+ 
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, inner_patches=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.inner_patches = inner_patches
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        in_chans = 6
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        conv_size = [size // inner_patches for size in patch_size]
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=conv_size, stride=conv_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        patches_resolution = (H // self.patch_size[0], W // self.patch_size[1])
+        num_patches = patches_resolution[0] * patches_resolution[1]
+        x = self.proj(x).view(
+            B, -1, 
+            patches_resolution[0], self.inner_patches, 
+            patches_resolution[1], self.inner_patches, 
+        ).permute(0, 2, 4, 3, 5, 1).reshape(B, num_patches, self.inner_patches, self.inner_patches, -1)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+class PatchMerge(nn.Module):
+    def __init__(self, dim, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(dim * 4)
+        self.reduction = nn.Linear(dim * 4, dim * 2, bias=False)
+    
+    def forward(self, x):
+        x0 = x[..., 0::2, 0::2, :] 
+        x1 = x[..., 1::2, 0::2, :] 
+        x2 = x[..., 0::2, 1::2, :] 
+        x3 = x[..., 1::2, 1::2, :]
+
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+        x = self.norm(x)
+        x = self.reduction(x)
+        return x
+
+
+class HiViT(BaseBackbone):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000,
+                 embed_dim=512, depths=[4, 4, 20], num_heads=8, stem_mlp_ratio=3., mlp_ratio=4., 
+                 qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.0,
+                 norm_layer=nn.LayerNorm, ape=True, rpe=True, patch_norm=True, use_checkpoint=False, 
+                 **kwargs):
+        super().__init__()
+        self.num_layers = len(depths)
+        self.ape = ape
+        self.rpe = rpe
+        self.patch_norm = patch_norm
+        self.num_features = self.embed_dim = embed_dim
+        self.mlp_ratio = mlp_ratio
+        self.use_checkpoint = use_checkpoint
+        self.num_main_blocks = depths[-1]
+
+        embed_dim = embed_dim // 2 ** (self.num_layers - 1)
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        Hp, Wp = self.patch_embed.patches_resolution
+        assert Hp == Wp
+
+        # absolute position embedding
+        if ape:
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, self.num_features)
+            )
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+        if rpe:
+            coords_h = torch.arange(Hp)
+            coords_w = torch.arange(Wp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w])) 
+            coords_flatten = torch.flatten(coords, 1) 
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] 
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous() 
+            relative_coords[:, :, 0] += Hp - 1 
+            relative_coords[:, :, 1] += Wp - 1
+            relative_coords[:, :, 0] *= 2 * Wp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer("relative_position_index", relative_position_index)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = iter(x.item() for x in torch.linspace(0, drop_path_rate, sum(depths) + sum(depths[:-1])))  # stochastic depth decay rule
+
+        # build blocks
+        self.blocks = nn.ModuleList()
+        for stage_depth in depths:
+            is_main_stage = embed_dim == self.num_features
+            nhead = num_heads if is_main_stage else 0
+            ratio = mlp_ratio if is_main_stage else stem_mlp_ratio
+            # every block not in main stage include two mlp blocks
+            stage_depth = stage_depth if is_main_stage else stage_depth * 2
+            for i in range(stage_depth):
+                self.blocks.append(
+                    BlockWithRPE(
+                        Hp, embed_dim, nhead, ratio, qkv_bias, qk_scale, 
+                        drop=drop_rate, attn_drop=attn_drop_rate, drop_path=next(dpr), 
+                        rpe=rpe, norm_layer=norm_layer,
+                    )
+                )
+            if not is_main_stage:
+                self.blocks.append(
+                    PatchMerge(embed_dim, norm_layer)
+                )
+                embed_dim *= 2
+        self.norm_ = norm_layer(embed_dim)
+        # self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+def _create_vision_transformer(pretrained=False, default_cfg=None, **kwargs):
+    
+    model = HiViT(**kwargs)
+
+    if pretrained:
+            checkpoint = torch.load(pretrained, map_location="cpu")
+            missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+            print(missing_keys, unexpected_keys)
+            print('Load pretrained model from: ' + pretrained)
+
+    return model
+
+
+def hivit_base(pretrained=False, **kwargs):
+    model_kwargs = dict(
+        embed_dim=512, depths=[2, 2, 20], num_heads=8, stem_mlp_ratio=3., mlp_ratio=4., 
+        rpe=False, norm_layer=partial(nn.LayerNorm, eps=1e-6),**kwargs
+    )
+    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
+    
+    return model
+
+def hivit_small(pretrained=False, **kwargs):
+    model_kwargs = dict(
+        embed_dim=384, depths=[2, 2, 20], num_heads=6, stem_mlp_ratio=3., mlp_ratio=4., 
+        rpe=False, norm_layer=partial(nn.LayerNorm, eps=1e-6),**kwargs
+    )
+    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
+    
+    return model

lib/models/aqatrack/itpn.py

@@ -0,0 +1,430 @@
+# --------------------------------------------------------
+# BEIT: BERT Pre-Training of Image Transformers (https://arxiv.org/abs/2106.08254)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# By Yunjie Tian
+# Based on timm and DeiT code bases
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# https://github.com/facebookresearch/deit/
+# https://github.com/facebookresearch/dino
+# --------------------------------------------------------'
+from functools import partial
+
+import math
+import torch
+import torch.nn as nn
+from timm.models.registry import register_model
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.vision_transformer import DropPath, Mlp, trunc_normal_
+from timm.models.layers import to_2tuple
+from lib.models.aqatrack.base_backbone import BaseBackbone
+from lib.utils.utils import combine_tokens, recover_tokens
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
+        **kwargs
+    }
+
+
+class Attention(nn.Module):
+    def __init__(self, input_size, dim, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0., rpe=True):
+        super().__init__()
+        self.input_size = input_size
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * input_size - 1) * (2 * input_size - 1), num_heads)
+        ) if rpe else None
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, rpe_index=None, mask=None):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        if rpe_index is not None:
+            S = int(math.sqrt(rpe_index.size(-1)))
+            relative_position_bias = self.relative_position_bias_table[rpe_index].view(-1, S, S, self.num_heads)
+            relative_position_bias = relative_position_bias.permute(0, 3, 1, 2).contiguous()
+            attn = attn + relative_position_bias
+        if mask is not None:
+            mask = mask.bool()
+            attn = attn.masked_fill(~mask[:, None, None, :], float("-inf"))
+        attn = attn.float().clamp(min=torch.finfo(torch.float32).min, max=torch.finfo(torch.float32).max)
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class BlockWithRPE(nn.Module):
+    def __init__(self, input_size, dim, num_heads=0., mlp_ratio=4., qkv_bias=True, qk_scale=None, init_values=None,
+                 drop=0., attn_drop=0., drop_path=0., rpe=True, act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+
+        with_attn = num_heads > 0.
+
+        self.norm1 = norm_layer(dim) if with_attn else None
+        self.attn = Attention(
+            input_size, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, rpe=rpe,
+        ) if with_attn else None
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values is not None and init_values > 0:
+            if self.attn is not None:
+                self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            else:
+                self.gamma_1 = None
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, rpe_index=None, mask=None):
+        if self.gamma_2 is None:
+            if self.attn is not None:
+                x = x + self.drop_path(self.attn(self.norm1(x), rpe_index, mask))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.attn is not None:
+                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rpe_index, mask))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, inner_patches=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.inner_patches = inner_patches
+        self.patches_resolution = self.patch_shape = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        in_chans = 6
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        conv_size = [size // inner_patches for size in patch_size]
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=conv_size, stride=conv_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        patches_resolution = (H // self.patch_size[0], W // self.patch_size[1])
+        num_patches = patches_resolution[0] * patches_resolution[1]
+        x = self.proj(x).view(
+            B, -1,
+            patches_resolution[0], self.inner_patches,
+            patches_resolution[1], self.inner_patches,
+        ).permute(0, 2, 4, 3, 5, 1).reshape(B, num_patches, self.inner_patches, self.inner_patches, -1)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+# the spatial size are split into 4 patches (then downsample 2x)
+# concat them and then reduce them to be of 2x channels.
+#
+class PatchMerge(nn.Module):
+    def __init__(self, dim, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(dim * 4)
+        self.reduction = nn.Linear(dim * 4, dim * 2, bias=False)
+
+    def forward(self, x):
+        x0 = x[..., 0::2, 0::2, :]
+        x1 = x[..., 1::2, 0::2, :]
+        x2 = x[..., 0::2, 1::2, :]
+        x3 = x[..., 1::2, 1::2, :]
+
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+        x = self.norm(x)
+        x = self.reduction(x)
+        return x
+
+
+# PatchSplit is for upsample
+# used in FPN for downstream tasks such as detection/segmentation
+class PatchSplit(nn.Module):
+    def __init__(self, dim, fpn_dim, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(dim)
+        self.reduction = nn.Linear(dim, fpn_dim * 4, bias=False)
+        self.fpn_dim = fpn_dim
+
+    def forward(self, x):
+        B, N, H, W, C = x.shape
+        x = self.norm(x)
+        x = self.reduction(x)
+        x = x.reshape(
+            B, N, H, W, 2, 2, self.fpn_dim
+        ).permute(0, 1, 2, 4, 3, 5, 6).reshape(
+            B, N, 2 * H, 2 * W, self.fpn_dim
+        )
+        return x
+
+
+class iTPN(BaseBackbone):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, fpn_dim=256, fpn_depth=2,
+                 embed_dim=256, mlp_depth1=3, mlp_depth2=3, depth=24, num_heads=8, bridge_mlp_ratio=3., mlp_ratio=4.,
+                 qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.0, init_values=None,
+                 norm_layer=nn.LayerNorm, ape=True, rpe=False, patch_norm=True, use_checkpoint=False,
+                 num_outs=3, **kwargs):
+        super().__init__()
+        assert num_outs in [1, 2, 3, 4, 5]
+        self.num_classes = num_classes
+        self.ape = ape
+        self.rpe = rpe
+        self.patch_norm = patch_norm
+        self.num_features = self.embed_dim = embed_dim
+        self.mlp_ratio = mlp_ratio
+        self.use_checkpoint = use_checkpoint
+        self.num_outs = num_outs
+        self.num_main_blocks = depth
+        self.fpn_dim = fpn_dim
+        self.mlp_depth1 = mlp_depth1
+        self.mlp_depth2 = mlp_depth2
+        self.depth = depth
+
+        mlvl_dims = {'4': embed_dim // 4, '8': embed_dim // 2, '16': embed_dim}
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=mlvl_dims['4'],
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        Hp, Wp = self.patch_embed.patches_resolution
+        assert Hp == Wp
+
+        # absolute position embedding
+        if ape:
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, self.num_features)
+            )
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+        if rpe:
+            coords_h = torch.arange(Hp)
+            coords_w = torch.arange(Wp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
+            coords_flatten = torch.flatten(coords, 1)
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+            relative_coords[:, :, 0] += Hp - 1
+            relative_coords[:, :, 1] += Wp - 1
+            relative_coords[:, :, 0] *= 2 * Wp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer("relative_position_index", relative_position_index)
+
+            # Hp_z, Hp_z
+            coords_h = torch.arange(Hp)
+            coords_w = torch.arange(Wp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
+            coords_flatten = torch.flatten(coords, 1)
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+            relative_coords[:, :, 0] += Hp - 1
+            relative_coords[:, :, 1] += Wp - 1
+            relative_coords[:, :, 0] *= 2 * Wp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer("relative_position_index", relative_position_index)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = iter(x.item() for x in torch.linspace(0, drop_path_rate, mlp_depth1 + mlp_depth2 + depth))
+        self.blocks = nn.ModuleList()
+
+        self.blocks.extend([
+            BlockWithRPE(
+                Hp, mlvl_dims['4'], 0, bridge_mlp_ratio, qkv_bias, qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=next(dpr),
+                rpe=rpe, norm_layer=norm_layer, init_values=init_values
+            ) for _ in range(mlp_depth1)]
+        )
+        self.blocks.append(PatchMerge(mlvl_dims['4'], norm_layer))
+        self.blocks.extend([
+            BlockWithRPE(
+                Hp, mlvl_dims['8'], 0, bridge_mlp_ratio, qkv_bias, qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=next(dpr),
+                rpe=rpe, norm_layer=norm_layer, init_values=init_values
+            ) for _ in range(mlp_depth2)]
+        )
+        self.blocks.append(PatchMerge(mlvl_dims['8'], norm_layer))
+        self.blocks.extend([
+            BlockWithRPE(
+                Hp, mlvl_dims['16'], num_heads, mlp_ratio, qkv_bias, qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=next(dpr),
+                rpe=rpe, norm_layer=norm_layer, init_values=init_values
+            ) for _ in range(depth)]
+        )
+
+        self.norm_ = norm_layer(embed_dim)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def interpolate_pos_encoding(self, x, h, w):
+        npatch = x.shape[1]
+        N = self.absolute_pos_embed.shape[1]
+        if npatch == N and w == h:
+            return self.absolute_pos_embed
+        patch_pos_embed = self.absolute_pos_embed
+        dim = x.shape[-1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w + 0.1, h + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(h0 / math.sqrt(N), w0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return patch_pos_embed
+
+    def get_num_layers(self):
+        return self.mlp_depth1 + self.mlp_depth2 + self.depth
+
+    def forward_features(self, z,x,soft_token_template_mask):
+        B, C, H, W = x.shape
+        Hp, Wp = H // self.patch_embed.patch_size[0], W // self.patch_embed.patch_size[1]
+
+        _, _, H_z, W_z = z.shape
+        Hp_z, Wp_z = H_z // self.patch_embed.patch_size[0], W_z // self.patch_embed.patch_size[1]
+
+        z = self.patch_embed(z)
+        x = self.patch_embed(x)
+
+        for blk in self.blocks[:-self.num_main_blocks]:
+            x =  blk(x)
+            z = blk(z)
+
+        x = x[..., 0, 0, :]
+        z = z[..., 0, 0, :]
+
+        x += self.pos_embed_x
+        z += self.pos_embed_z
+
+        token_type_search = self.token_type_search.expand(B, x.shape[1], -1)
+        token_type_template_bg = self.token_type_template_bg.expand(B, z.shape[1], -1)
+        token_type_template_fg = self.token_type_template_fg.expand(B, z.shape[1], -1)
+
+        x += token_type_search
+
+        token_type_template_fg = soft_token_template_mask*token_type_template_fg
+        token_type_template_bg = (1- soft_token_template_mask) * token_type_template_bg
+        z += token_type_template_fg + token_type_template_bg
+
+        x = combine_tokens(z, x, mode=self.cat_mode)
+        if self.add_cls_token:
+            cls_tokens = self.cls_token.expand(B, -1, -1)
+            # cls_tokens = cls_tokens + self.cls_pos_embed
+            x = torch.cat([cls_tokens, x], dim=1)
+
+        x = self.pos_drop(x)
+
+        for blk in self.blocks[-self.num_main_blocks:]:
+            x = blk(x)
+
+        lens_z = self.pos_embed_z.shape[1]
+        lens_x = self.pos_embed_x.shape[1]
+        x = recover_tokens(x, lens_z, lens_x, mode=self.cat_mode)
+
+        aux_dict = {"attn": None}
+        x = self.norm_(x)
+
+        return x,aux_dict
+
+
+    def forward(self, z,x,soft_token_template_mask, **kwargs):
+        x,aux_dict = self.forward_features(z,x,soft_token_template_mask)
+        return x,aux_dict
+
+
+@register_model
+def itpn_base_3324_patch16_224(pretrained="", **kwargs):
+    model = iTPN(
+        patch_size=16, embed_dim=512, mlp_depth1=3, mlp_depth2=3, depth=24, num_heads=8, bridge_mlp_ratio=3.,
+        mlp_ratio=4, num_outs=1, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        checkpoint = torch.load(
+            pretrained, map_location="cpu"
+        )
+        ###### naive method
+        # missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+
+        #### process patch_embed.proj
+        normal_weights = {}
+        # special_weights = {}
+        for k, v in checkpoint.items():
+            if "patch_embed.proj.weight" in k:
+                normal_weights[k] = torch.cat([v, v], dim=1)
+            else:
+                normal_weights[k] = v
+
+        missing_keys, unexpected_keys = model.load_state_dict(normal_weights, strict=False)
+
+    return model
+
+
+@register_model
+def itpn_large_2240_patch16_224(pretrained=False, **kwargs):
+    model = iTPN(
+        patch_size=16, embed_dim=768, mlp_depth1=2, mlp_depth2=2, depth=40, num_heads=12, bridge_mlp_ratio=3.,
+        mlp_ratio=4, num_outs=1, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        checkpoint = torch.load(
+            kwargs["init_ckpt"], map_location="cpu"
+        )
+        model.load_state_dict(checkpoint["model"])
+    return model

lib/models/aqatrack/utils.py

@@ -0,0 +1,99 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+def combine_tokens(template_tokens, search_tokens, mode='direct', return_res=False):
+    # [B, HW, C]
+    len_t = template_tokens.shape[1]
+    len_s = search_tokens.shape[1]
+
+    if mode == 'direct':
+        merged_feature = torch.cat((template_tokens, search_tokens), dim=1)
+    elif mode == 'template_central':
+        central_pivot = len_s // 2
+        first_half = search_tokens[:, :central_pivot, :]
+        second_half = search_tokens[:, central_pivot:, :]
+        merged_feature = torch.cat((first_half, template_tokens, second_half), dim=1)
+    elif mode == 'partition':
+        feat_size_s = int(math.sqrt(len_s))
+        feat_size_t = int(math.sqrt(len_t))
+        window_size = math.ceil(feat_size_t / 2.)
+        # pad feature maps to multiples of window size
+        B, _, C = template_tokens.shape
+        H = W = feat_size_t
+        template_tokens = template_tokens.view(B, H, W, C)
+        pad_l = pad_b = pad_r = 0
+        # pad_r = (window_size - W % window_size) % window_size
+        pad_t = (window_size - H % window_size) % window_size
+        template_tokens = F.pad(template_tokens, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = template_tokens.shape
+        template_tokens = template_tokens.view(B, Hp // window_size, window_size, W, C)
+        template_tokens = torch.cat([template_tokens[:, 0, ...], template_tokens[:, 1, ...]], dim=2)
+        _, Hc, Wc, _ = template_tokens.shape
+        template_tokens = template_tokens.view(B, -1, C)
+        merged_feature = torch.cat([template_tokens, search_tokens], dim=1)
+
+        # calculate new h and w, which may be useful for SwinT or others
+        merged_h, merged_w = feat_size_s + Hc, feat_size_s
+        if return_res:
+            return merged_feature, merged_h, merged_w
+
+    else:
+        raise NotImplementedError
+
+    return merged_feature
+
+
+def recover_tokens(merged_tokens, len_template_token, len_search_token, mode='direct'):
+    if mode == 'direct':
+        recovered_tokens = merged_tokens
+    elif mode == 'template_central':
+        central_pivot = len_search_token // 2
+        len_remain = len_search_token - central_pivot
+        len_half_and_t = central_pivot + len_template_token
+
+        first_half = merged_tokens[:, :central_pivot, :]
+        second_half = merged_tokens[:, -len_remain:, :]
+        template_tokens = merged_tokens[:, central_pivot:len_half_and_t, :]
+
+        recovered_tokens = torch.cat((template_tokens, first_half, second_half), dim=1)
+    elif mode == 'partition':
+        recovered_tokens = merged_tokens
+    else:
+        raise NotImplementedError
+
+    return recovered_tokens
+
+
+def window_partition(x, window_size: int):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size: int, H: int, W: int):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    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)
+    return x

lib/models/atctrack/__init__.py

@@ -0,0 +1 @@
+from .atctrack import build_atctrack

lib/models/atctrack/atctrack.py

@@ -0,0 +1,1051 @@
+"""
+ATCTrack  Model
+"""
+import os
+import re
+import base64
+from io import BytesIO
+
+import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+
+from lib.utils.misc import NestedTensor
+
+# from .language_model import build_bert
+from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy, box_xyxy_to_cxcywh, box_iou
+### aqatrack
+from lib.models.aqatrack.hivit import hivit_small, hivit_base
+from lib.models.aqatrack.itpn import itpn_base_3324_patch16_224
+from lib.models.aqatrack.fast_itpn import fast_itpn_base_3324_patch16_224,fast_itpn_large_2240_patch16_256
+
+from lib.models.transformers.transformer import build_rgb_det_decoder
+from lib.models.layers.transformer_dec import build_transformer_dec,build_transformer_dec_with_mask
+
+from torch.nn.modules.transformer import _get_clones
+from lib.models.layers.head import build_box_head
+
+import torch.nn.functional as F
+from lib.models.layers.frozen_bn import FrozenBatchNorm2d
+from transformers import BertTokenizer, BertModel, RobertaModel, RobertaTokenizerFast, AutoTokenizer, AutoProcessor
+from PIL import Image, ImageDraw
+from lib.models.transformers import build_decoder, VisionLanguageFusionModule, PositionEmbeddingSine1D,build_text_prompt_decoder
+
+
+TARGET_STATE_TOKEN = "<TARGET_STATE>"
+
+
+def _project_root():
+    return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../.."))
+
+
+def _resolve_project_path(path):
+    if not path or os.path.isabs(path):
+        return path
+    candidate = os.path.abspath(os.path.join(_project_root(), path))
+    if os.path.exists(candidate) or path.startswith((".", "..", "checkpoint", "resource")):
+        return candidate
+    return path
+
+
+def _load_qwen_target_state_model(model_path):
+    try:
+        from transformers import AutoModelForImageTextToText
+        model_cls = AutoModelForImageTextToText
+    except ImportError:
+        from transformers import AutoModelForCausalLM
+        model_cls = AutoModelForCausalLM
+
+    try:
+        return model_cls.from_pretrained(model_path, trust_remote_code=True)
+    except ValueError as exc:
+        raise RuntimeError(
+            "Cannot load Qwen target-state model. The current transformers package "
+            "does not recognize this Qwen architecture. Upgrade transformers in the "
+            "training environment before enabling MODEL.TARGET_STATE."
+        ) from exc
+
+
+class QwenTargetStateEncoder(nn.Module):
+    def __init__(self, cfg, tracker_dim):
+        super().__init__()
+        ts_cfg = cfg.MODEL.TARGET_STATE
+        self.model_path = _resolve_project_path(os.environ.get("QWEN_MODEL_PATH", ts_cfg.MODEL_PATH))
+        self.token = getattr(ts_cfg, "TOKEN", TARGET_STATE_TOKEN)
+        self.prompt_template = getattr(ts_cfg, "PROMPT_TEMPLATE", "default")
+        self.train_token_embedding = getattr(ts_cfg, "TRAIN_TOKEN_EMBEDDING", False)
+        self.freeze_qwen = getattr(ts_cfg, "FREEZE_QWEN", True)
+        self.use_lora = getattr(ts_cfg, "USE_LORA", False)
+        self.lora_r = getattr(ts_cfg, "LORA_R", 8)
+        self.lora_alpha = getattr(ts_cfg, "LORA_ALPHA", 16)
+        self.lora_dropout = getattr(ts_cfg, "LORA_DROPOUT", 0.05)
+        self.lora_target_modules = getattr(ts_cfg, "LORA_TARGET_MODULES", [
+            "in_proj_qkv", "out_proj", "in_proj_z", "in_proj_b", "in_proj_a",
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ])
+        teacher_enable_env = os.environ.get("QWEN_TEACHER_ENABLE")
+        if teacher_enable_env is None:
+            self.teacher_enable = bool(getattr(ts_cfg, "TEACHER_ENABLE", False))
+        else:
+            self.teacher_enable = teacher_enable_env.strip().lower() in ("1", "true", "yes", "on")
+        self.teacher_model = os.environ.get("QWEN_TEACHER_MODEL", getattr(ts_cfg, "TEACHER_MODEL", "qwen3.5"))
+        self.teacher_base_url = os.environ.get("QWEN_TEACHER_BASE_URL", getattr(ts_cfg, "TEACHER_BASE_URL", "http://127.0.0.1:8001/v1"))
+        self.teacher_api_key = os.environ.get("QWEN_TEACHER_API_KEY", getattr(ts_cfg, "TEACHER_API_KEY", "sk-no-key-required"))
+        self.teacher_client = None
+
+        self.processor = AutoProcessor.from_pretrained(self.model_path, trust_remote_code=True)
+        self.tokenizer = getattr(self.processor, "tokenizer", None)
+        if self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "left"
+        self.qwen = _load_qwen_target_state_model(self.model_path)
+
+        self.target_state_special_tokens = ["<answer>", "</answer>", "<state_token>", "</state_token>", self.token]
+        special_tokens = {"additional_special_tokens": self.target_state_special_tokens}
+        num_added = self.tokenizer.add_special_tokens(special_tokens)
+        if num_added > 0:
+            self.qwen.resize_token_embeddings(len(self.tokenizer))
+        self.target_token_id = self.tokenizer.convert_tokens_to_ids(self.token)
+        self._embedding_grad_hook = None
+
+        qwen_hidden_dim = self.qwen.config.text_config.hidden_size if hasattr(self.qwen.config, "text_config") else self.qwen.config.hidden_size
+        self.projector = nn.Sequential(
+            nn.Linear(qwen_hidden_dim, tracker_dim),
+            nn.LayerNorm(tracker_dim),
+            nn.GELU(),
+            nn.Linear(tracker_dim, tracker_dim),
+        )
+        self.film = nn.Linear(tracker_dim, tracker_dim * 2)
+        self.film_gate = nn.Parameter(torch.zeros(1))
+
+        if self.freeze_qwen:
+            for p in self.qwen.parameters():
+                p.requires_grad = False
+        if self.use_lora:
+            self._enable_qwen_lora()
+        self.configure_token_embedding_training(self.train_token_embedding)
+
+    def _enable_qwen_lora(self):
+        try:
+            from peft import LoraConfig, get_peft_model
+        except ImportError as exc:
+            raise RuntimeError("MODEL.TARGET_STATE.USE_LORA=True requires the peft package.") from exc
+
+        target_modules = self.lora_target_modules
+        if isinstance(target_modules, str):
+            target_modules = [item.strip() for item in target_modules.split(",") if item.strip()]
+        config = LoraConfig(
+            r=self.lora_r,
+            lora_alpha=self.lora_alpha,
+            target_modules=list(target_modules),
+            lora_dropout=self.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        self.qwen = get_peft_model(self.qwen, config)
+
+    def configure_token_embedding_training(self, enabled):
+        embedding = self.qwen.get_input_embeddings()
+        embedding.weight.requires_grad = bool(enabled)
+        if self._embedding_grad_hook is not None:
+            self._embedding_grad_hook.remove()
+            self._embedding_grad_hook = None
+        if enabled:
+            train_token_ids = torch.tensor([self.target_token_id], dtype=torch.long)
+
+            def mask_embedding_grad(grad):
+                token_ids = train_token_ids.to(grad.device)
+                mask = torch.zeros((grad.shape[0],), device=grad.device, dtype=grad.dtype)
+                mask.index_fill_(0, token_ids, 1)
+                return grad * mask.view(-1, 1)
+
+            self._embedding_grad_hook = embedding.weight.register_hook(mask_embedding_grad)
+
+    def _qwen_forward_with_target_embedding(self, tokenized, labels=None):
+        return self.qwen(**tokenized, labels=labels, output_hidden_states=True, use_cache=False)
+
+    @staticmethod
+    def _tensor_batch_to_pil(images, boxes=None):
+        mean = images.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+        std = images.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+        images = (images.detach().float() * std + mean).clamp(0, 1)
+        images = (images * 255).byte().permute(0, 2, 3, 1).cpu().numpy()
+        pil_images = [Image.fromarray(image) for image in images]
+
+        if boxes is None:
+            return pil_images
+
+        boxes = boxes.detach().float().cpu()
+        for image, box in zip(pil_images, boxes):
+            draw = ImageDraw.Draw(image)
+            x, y, w, h = box.tolist()
+            if max(abs(x), abs(y), abs(w), abs(h)) <= 2.0:
+                img_w, img_h = image.size
+                x, w = x * img_w, w * img_w
+                y, h = y * img_h, h * img_h
+            x1 = max(0.0, min(float(image.size[0] - 1), x))
+            y1 = max(0.0, min(float(image.size[1] - 1), y))
+            x2 = max(0.0, min(float(image.size[0] - 1), x + w))
+            y2 = max(0.0, min(float(image.size[1] - 1), y + h))
+            if x2 > x1 and y2 > y1:
+                line_width = max(2, round(min(image.size) / 80))
+                draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=line_width)
+        return pil_images
+
+    def _build_prompt(self, caption, object_name=None):
+        caption = caption if caption else "the target object"
+        object_name = object_name if object_name else caption
+        return (
+            f"Role: {object_name} tracking update judge and target-state token generator.\n\n"
+            "Task: Compare the targets inside the provided bboxes in Frame 1 (Original) "
+            "and Frame 2 (New). Decide whether Frame 2 should update the tracking template, "
+            "and generate a target-state token for the tracking model.\n\n"
+            "Reject update for full occlusion, out of view, too small target, severe blur/clipping, "
+            "wrong bbox, distractor, uncertain identity, or no meaningful target appearance change.\n\n"
+            "Accept update only if Frame 2 contains the same target as Frame 1, the bbox is reliable, "
+            "the target is clear, and the appearance change is useful for future tracking.\n\n"
+            "The target-state token should summarize the current target condition for the tracking model. "
+            "It should encode whether the candidate is reliable, whether the target identity is consistent, "
+            "and whether the current appearance is useful or risky for tracking.\n\n"
+            "Frame 1 (Original) is the first image. Frame 2 (New candidate/search crop) is the second image.\n\n"
+            "Output exactly one answer XML tag containing yes or no, immediately followed by one "
+            "state_token XML tag containing the special target-state token. Do not output any extra text."
+        )
+
+    def _build_teacher_prompt(self, caption, object_name=None):
+        caption = caption if caption else "the target object"
+        object_name = object_name if object_name else caption
+        return (
+            f"Role: {object_name} tracking update judge.\n"
+            "Task: Compare the targets inside the provided bboxes in Frame 1 (Original) and Frame 2 (New), "
+            "and decide whether Frame 2 should update the tracking template.\n\n"
+            "Reject update for full occlusion, out of view, too small target, severe blur/clipping, wrong bbox, "
+            "distractor, uncertain identity, or no meaningful target appearance change.\n"
+            "Accept only if Frame 2 is the same target, bbox is reliable, target is clear, and appearance change is useful.\n\n"
+            "Output exactly: <answer>yes/no</answer>"
+        )
+
+    @staticmethod
+    def _pil_to_base64_jpeg(image):
+        buffer = BytesIO()
+        image.save(buffer, format="JPEG")
+        return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    def _get_teacher_client(self):
+        if self.teacher_client is None:
+            try:
+                from openai import OpenAI
+            except ImportError as exc:
+                raise RuntimeError("Teacher update judge requires the openai package.") from exc
+            self.teacher_client = OpenAI(base_url=self.teacher_base_url, api_key=self.teacher_api_key, timeout=5.0)
+        return self.teacher_client
+
+    def _query_teacher_decisions(self, prompts, template_pils, search_pils):
+        if not self.teacher_enable:
+            return None, None
+        client = self._get_teacher_client()
+        decisions = []
+        responses = []
+        for prompt, template_pil, search_pil in zip(prompts, template_pils, search_pils):
+            base64_image1 = self._pil_to_base64_jpeg(template_pil)
+            base64_image2 = self._pil_to_base64_jpeg(search_pil)
+            messages = [{
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image1}"}},
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image2}"}},
+                    {"type": "text", "text": prompt},
+                ],
+            }]
+            try:
+                chat_response = client.chat.completions.create(
+                    model=self.teacher_model,
+                    messages=messages,
+                    max_tokens=32,
+                    temperature=0.1,
+                    top_p=1.0,
+                    presence_penalty=0.0,
+                    extra_body={
+                        "top_k": 1,
+                        "chat_template_kwargs": {"enable_thinking": False},
+                        "guided_regex": r"<answer>(yes|no)</answer>",
+                    },
+                )
+                content = chat_response.choices[0].message.content
+                match = re.findall(r"<answer>\s*(yes|no)\s*</answer>", content, flags=re.IGNORECASE)
+                if not match:
+                    decisions.append(None)
+                else:
+                    decisions.append(match[-1].lower() == "yes")
+                responses.append(content)
+            except Exception as exc:
+                print(f"Qwen teacher update judge failed: {exc}")
+                decisions.append(None)
+                responses.append(None)
+        return decisions, responses
+
+    @staticmethod
+    def _parse_update_decisions(decoded_outputs):
+        decisions = []
+        for text in decoded_outputs:
+            text_l = text.lower()
+            answer_start = text_l.rfind("<answer>")
+            answer_end = text_l.find("</answer>", answer_start + len("<answer>")) if answer_start >= 0 else -1
+            answer = text_l[answer_start + len("<answer>"):answer_end].strip() if answer_start >= 0 and answer_end >= 0 else text_l
+            answer = answer.replace("<|im_end|>", " ").replace("<|endoftext|>", " ")
+            tokens = answer.replace("<", " ").replace(">", " ").replace("/", " ").split()
+            if "yes" in tokens and "no" not in tokens:
+                decisions.append(True)
+            elif "no" in tokens:
+                decisions.append(False)
+            else:
+                decisions.append(False)
+        return decisions
+
+    def _apply_qwen_chat_template(self, message):
+        try:
+            return self.processor.apply_chat_template(
+                message,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+        except TypeError:
+            return self.processor.apply_chat_template(
+                message,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+
+    def _target_state_answer_sequences(self):
+        outputs = [
+            f"<answer>yes</answer><state_token>{self.token}</state_token>",
+            f"<answer>no</answer><state_token>{self.token}</state_token>",
+        ]
+        return [self.tokenizer(text, add_special_tokens=False).input_ids for text in outputs]
+
+    def _qwen_generate(self, **kwargs):
+        if self.training and hasattr(self.qwen, "get_base_model"):
+            base_model = self.qwen.get_base_model()
+            unwrapped = getattr(base_model.generate, "__wrapped__", None)
+            if unwrapped is not None:
+                with self.qwen._enable_peft_forward_hooks(**kwargs):
+                    peft_args = getattr(self.qwen, "special_peft_forward_args", set())
+                    clean_kwargs = {k: v for k, v in kwargs.items() if k not in peft_args}
+                    return unwrapped(base_model, **clean_kwargs)
+
+        generate_fn = self.qwen.generate
+        if self.training:
+            unwrapped = getattr(generate_fn, "__wrapped__", None)
+            if unwrapped is not None:
+                return unwrapped(self.qwen, **kwargs)
+        return generate_fn(**kwargs)
+
+    def _qwen_generation_kwargs(self, prompt_len=None):
+        eos_token_ids = []
+        for token in ("<|im_end|>", "<|endoftext|>"):
+            token_id = self.tokenizer.convert_tokens_to_ids(token)
+            if isinstance(token_id, int) and token_id >= 0 and token_id != self.tokenizer.unk_token_id:
+                eos_token_ids.append(token_id)
+        if self.tokenizer.eos_token_id is not None:
+            eos_token_ids.append(self.tokenizer.eos_token_id)
+        eos_token_ids = list(dict.fromkeys(eos_token_ids))
+        kwargs = {
+            "max_new_tokens": 16,
+            "do_sample": False,
+            "num_beams": 1,
+            "repetition_penalty": 1.0,
+            "eos_token_id": eos_token_ids or self.tokenizer.eos_token_id,
+            "pad_token_id": self.tokenizer.pad_token_id,
+        }
+        if prompt_len is not None:
+            answer_sequences = self._target_state_answer_sequences()
+            stop_ids = eos_token_ids or [self.tokenizer.eos_token_id]
+
+            def prefix_allowed_tokens_fn(batch_id, input_ids):
+                suffix = input_ids[prompt_len:].tolist()
+                allowed = []
+                for sequence in answer_sequences:
+                    if len(suffix) <= len(sequence) and suffix == sequence[:len(suffix)]:
+                        if len(suffix) == len(sequence):
+                            allowed.extend(stop_ids)
+                        else:
+                            allowed.append(sequence[len(suffix)])
+                return list(dict.fromkeys(allowed)) or stop_ids
+
+            kwargs["prefix_allowed_tokens_fn"] = prefix_allowed_tokens_fn
+        return kwargs
+
+    def _format_loss_from_generation_scores(self, scores, generated_suffix):
+        if scores is None or len(scores) == 0:
+            return generated_suffix.new_tensor(0.0, dtype=torch.float32)
+        num_steps = min(len(scores), generated_suffix.shape[1])
+        logits = torch.stack(scores[:num_steps], dim=1).float()
+        targets = generated_suffix[:, :num_steps]
+        return F.cross_entropy(
+            logits.reshape(-1, logits.shape[-1]),
+            targets.reshape(-1),
+            ignore_index=self.tokenizer.pad_token_id,
+        )
+
+    def _teacher_decision_loss(self, scores, teacher_decisions):
+        valid_items = [(idx, decision) for idx, decision in enumerate(teacher_decisions or []) if decision is not None]
+        if not valid_items or scores is None or len(scores) == 0:
+            device = scores[0].device if scores else self.qwen.get_input_embeddings().weight.device
+            return torch.tensor(0.0, device=device), None
+
+        yes_seq, no_seq = self._target_state_answer_sequences()
+        decision_step = next((i for i, (yes_id, no_id) in enumerate(zip(yes_seq, no_seq)) if yes_id != no_id), None)
+        if decision_step is None or decision_step >= len(scores):
+            return scores[0].new_tensor(0.0), None
+
+        batch_indices = torch.tensor([idx for idx, _ in valid_items], device=scores[decision_step].device, dtype=torch.long)
+        target_ids = torch.tensor(
+            [yes_seq[decision_step] if decision else no_seq[decision_step] for _, decision in valid_items],
+            device=scores[decision_step].device,
+            dtype=torch.long,
+        )
+        logits = scores[decision_step].float().index_select(0, batch_indices)
+        loss = F.cross_entropy(logits, target_ids)
+        labels = torch.full((len(teacher_decisions),), -1, device=scores[decision_step].device, dtype=torch.long)
+        labels[batch_indices] = torch.tensor([1 if decision else 0 for _, decision in valid_items], device=labels.device)
+        return loss, labels
+
+    def _target_hidden_from_generation(self, generation_hidden_states, generated_suffix):
+        target_mask = generated_suffix.eq(self.target_token_id)
+        if target_mask.any(dim=1).all():
+            target_pos = target_mask.float().argmax(dim=1)
+        else:
+            non_pad = generated_suffix.ne(self.tokenizer.pad_token_id)
+            target_pos = non_pad.sum(dim=1).clamp_min(1) - 1
+
+        hidden_steps = generation_hidden_states or []
+        if len(hidden_steps) == 0:
+            raise RuntimeError("Qwen generation did not return hidden states.")
+        h_targets = []
+        for batch_idx, pos in enumerate(target_pos.detach().cpu().tolist()):
+            # In cached generation, step t predicts generated token t. The hidden
+            # state for generated token k is available at step k + 1, when that
+            # token is fed back to predict the next token.
+            step = min(pos + 1, len(hidden_steps) - 1)
+            last_hidden = hidden_steps[step][-1]
+            h_targets.append(last_hidden[batch_idx, -1])
+        return torch.stack(h_targets, dim=0).float()
+
+    def forward(self, captions, template_images, search_images, template_boxes, search_boxes, device,
+                object_names=None, return_update_decision=False):
+        if object_names is None:
+            object_names = [None] * len(captions)
+        prompts = [self._build_prompt(caption, object_name) for caption, object_name in zip(captions, object_names)]
+        teacher_prompts = [self._build_teacher_prompt(caption, object_name) for caption, object_name in zip(captions, object_names)]
+        template_pils = self._tensor_batch_to_pil(template_images, template_boxes)
+        search_pils = self._tensor_batch_to_pil(search_images, search_boxes)
+        teacher_decisions, teacher_outputs = self._query_teacher_decisions(teacher_prompts, template_pils, search_pils)
+        messages = []
+        for prompt, template_pil, search_pil in zip(prompts, template_pils, search_pils):
+            messages.append([
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": template_pil},
+                        {"type": "image", "image": search_pil},
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ])
+        texts = [self._apply_qwen_chat_template(message) for message in messages]
+        images = [[template_pil, search_pil] for template_pil, search_pil in zip(template_pils, search_pils)]
+        tokenized = self.processor(text=texts, images=images, padding=True, return_tensors="pt").to(device)
+
+        generation = self._qwen_generate(
+            **tokenized,
+            **self._qwen_generation_kwargs(prompt_len=tokenized.input_ids.shape[1]),
+            return_dict_in_generate=True,
+            output_scores=True,
+            output_hidden_states=True,
+        )
+        generated_ids = generation.sequences
+        generated_suffix = generated_ids[:, tokenized.input_ids.shape[1]:]
+        decoded_outputs = self.tokenizer.batch_decode(generated_suffix, skip_special_tokens=False)
+        update_decisions = torch.tensor(
+            self._parse_update_decisions(decoded_outputs), device=device, dtype=torch.bool
+        )
+        qwen_format_loss = self._format_loss_from_generation_scores(generation.scores, generated_suffix)
+        qwen_teacher_loss, teacher_labels = self._teacher_decision_loss(generation.scores, teacher_decisions)
+        h_target = self._target_hidden_from_generation(generation.hidden_states, generated_suffix)
+        z_target = self.projector(h_target)
+        if return_update_decision:
+            return z_target, update_decisions, qwen_format_loss, qwen_teacher_loss, teacher_labels, teacher_outputs
+        return z_target
+
+    def modulate_feature(self, opt_feat, z_target):
+        gamma, beta = self.film(z_target).chunk(2, dim=-1)
+        gate = torch.tanh(self.film_gate)
+        gamma = gamma[:, :, None, None]
+        beta = beta[:, :, None, None]
+        return opt_feat * (1 + gate * gamma) + gate * beta
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1,
+         freeze_bn=False):
+    if freeze_bn:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            FrozenBatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+    else:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+class ConfidencePred(nn.Module):
+    def __init__(self):
+        super(ConfidencePred, self).__init__()
+        self.feat_sz = 24
+        self.stride = 1
+        self.img_sz = self.feat_sz * self.stride
+        freeze_bn = False
+
+        # CNN
+        self.conv1_ctr = conv(5, 16, freeze_bn=freeze_bn)
+        self.conv2_ctr = conv(16, 16 // 2, freeze_bn=freeze_bn)
+        self.conv3_ctr = conv(16 // 2, 16 // 4, freeze_bn=freeze_bn)
+        self.conv4_ctr = conv(16 // 4, 16 // 8, freeze_bn=freeze_bn)
+        self.conv5_ctr = nn.Conv2d(16 // 8, 1, kernel_size=1)
+
+        # 定义全连接层
+        self.fc1 = nn.Linear(256, 512)
+
+        ## cross attn 交互层
+        # self.multihead_attn = nn.MultiheadAttention(512, 4, dropout=0.1)
+        # # Implementation of Feedforward model
+        # self.dropout = nn.Dropout(0.1)
+        # self.norm1 = nn.LayerNorm(512)
+
+
+        self.fc2 = nn.Linear(512, 1)
+
+        # 定义激活函数
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x,xz_feature=None, gt_score_map=None):
+        """ Forward pass with input x. """
+
+        # ctr branch
+        x_ctr1 = self.conv1_ctr(x)
+        x_ctr2 = self.conv2_ctr(x_ctr1)
+        x_ctr3 = self.conv3_ctr(x_ctr2)
+        x_ctr4 = self.conv4_ctr(x_ctr3)
+        score_map_ctr = self.conv5_ctr(x_ctr4)
+
+        # 展平输入
+        x = score_map_ctr.flatten(1)
+        x = self.relu(self.fc1(x))
+
+        x = self.sigmoid(self.fc2(x))
+
+        return x
+
+class SubjectIndexPred(nn.Module):
+    def __init__(self,dim):
+        super(SubjectIndexPred, self).__init__()
+
+        # 定义全连接层
+        self.fc1 = nn.Linear(dim, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, 1)
+        self.sigmoid = nn.Sigmoid()
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x):
+        """ Forward pass with input x. """
+
+        # 全连接层前向传播
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.sigmoid(self.fc3(x))
+
+        return x
+
+
+class ATCTrack(nn.Module):
+    """ This is the base class for ATCTrack"""
+    def __init__(self, transformer,  box_head, tokenizer, text_encoder, aux_loss=False, head_type="CORNER",dim=512,cfg=None):
+        """ Initializes the model.
+        Parameters:
+            encoder: torch module of the encoder to be used. See encoder.py
+            decoder: torch module of the decoder architecture. See decoder.py
+        """
+        super().__init__()
+        self.backbone = transformer
+        self.box_head = box_head
+
+        self.aux_loss = aux_loss
+        self.head_type = head_type
+        if head_type == "CORNER" or head_type == "CENTER":
+            self.feat_sz_s = int(box_head.feat_sz)
+            self.feat_len_s = int(box_head.feat_sz ** 2)
+
+        if self.aux_loss:
+            self.box_head = _get_clones(self.box_head, 6)
+
+        self.dim = dim
+
+        self.query_len = 1
+        self.cls_prompts_pos = nn.Embedding(num_embeddings=self.query_len, embedding_dim=self.dim )  # pos for cur query
+        # self.cls_initial= nn.Embedding(num_embeddings=self.query_len, embedding_dim=self.dim )  # pos for cur query
+        self.confidence_pred = ConfidencePred()
+
+        ### visual temporal
+        self.visual_temporal_fusion = build_transformer_dec_with_mask(cfg, self.dim )
+        self.temporal_len = 4
+        self.dy_template_pos_embed = nn.Embedding(num_embeddings=self.temporal_len,
+                                                  embedding_dim=self.dim )  # pos for cur query
+
+        ## invlove_text
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.text_adj = nn.Sequential(
+            nn.Linear(768, self.dim , bias=True),
+            nn.LayerNorm(self.dim , eps=1e-12),
+            nn.Dropout(0.1),
+        )
+
+        self.language_adjust = build_transformer_dec(cfg, self.dim )
+        self.vl_fusion = VisionLanguageFusionModule(dim=self.dim , num_heads=8, attn_drop=0.1, proj_drop=0.1,
+                                                    num_vlfusion_layers=2,
+                                                    vl_input_type='separate')
+
+        self.text_pos = PositionEmbeddingSine1D(self.dim , normalize=True)
+
+        self.text_sub_idnex_classifier = SubjectIndexPred(self.dim)
+
+        self.use_target_state = getattr(cfg.MODEL.TARGET_STATE, "ENABLE", False) if hasattr(cfg.MODEL, "TARGET_STATE") else False
+        if self.use_target_state:
+            self.target_state_encoder = QwenTargetStateEncoder(cfg, self.dim)
+        else:
+            self.target_state_encoder = None
+
+    def forward_backbone(self, template, search, cls_token,soft_token_template_mask,x_pos):
+        # template b, 12, h,w
+        # search b,6,h,w
+        template = [template[:,:3],template[:,3:]]
+        soft_token_template_mask = [soft_token_template_mask[:, :64], soft_token_template_mask[:, 64:]]
+
+        x, token_type_infor = self.backbone.forward_features_pe(z=template, x=search, soft_token_template_mask =soft_token_template_mask)
+        x, aux_dict = self.backbone.forward_features_stage3(x, cls_token,x_pos)
+        return x, aux_dict
+
+    def forward(self, template: torch.Tensor,
+                search: torch.Tensor,
+                soft_token_template_mask=None,
+                exp_str=None,
+                exp_subject_mask=None,
+                target_state_exp_str=None,
+                target_state_template_bbox=None,
+                target_state_new_template_bbox=None,
+                target_state_object_name=None,
+                temporal_infor=[],
+                first_frame_flag=False,
+                training=True):
+
+        b0, num_search = template[0].shape[0], len(search)
+        z_target = None
+        target_state_update_decision = None
+        qwen_format_loss = None
+        qwen_teacher_loss = None
+        qwen_teacher_labels = None
+        qwen_teacher_outputs = None
+        target_state_captions = target_state_exp_str if target_state_exp_str is not None else exp_str
+
+        if training:
+            search = torch.cat(search, dim=0)
+            if self.use_target_state and target_state_captions and len(template) >= 3:
+                z_target, target_state_update_decision, qwen_format_loss, qwen_teacher_loss, qwen_teacher_labels, qwen_teacher_outputs = self.target_state_encoder(
+                    target_state_captions, template[-2], template[-1], target_state_template_bbox,
+                    target_state_new_template_bbox, search.device, object_names=target_state_object_name,
+                    return_update_decision=True
+                )
+                selector = target_state_update_decision.view(b0, 1, 1, 1)
+                dynamic_template = torch.where(selector, template[-1], template[-2])
+                dynamic_mask = torch.where(
+                    target_state_update_decision.view(b0, 1, 1),
+                    soft_token_template_mask[-1],
+                    soft_token_template_mask[-2],
+                )
+            else:
+                dynamic_template = template[1]
+                dynamic_mask = soft_token_template_mask[1]
+            template = torch.cat([template[0], dynamic_template], dim=1)
+            soft_token_template_mask = torch.cat([soft_token_template_mask[0], dynamic_mask], dim=1)
+            template_temporal = []
+            soft_token_template_mask_temporal = []
+            for _ in range(num_search):
+                template_temporal.append(template)
+                soft_token_template_mask_temporal.append(soft_token_template_mask)
+            template_temporal = torch.cat(template_temporal, dim=0)
+            soft_token_template_mask_temporal = torch.cat(soft_token_template_mask_temporal,dim=0)
+
+        else:
+            b0 = 1
+            if self.use_target_state and target_state_captions and len(template) >= 3:
+                z_target, target_state_update_decision, qwen_format_loss, qwen_teacher_loss, qwen_teacher_labels, qwen_teacher_outputs = self.target_state_encoder(
+                    target_state_captions, template[-2], template[-1], target_state_template_bbox,
+                    target_state_new_template_bbox, search.device, object_names=target_state_object_name,
+                    return_update_decision=True
+                )
+                dynamic_template = template[-1] if bool(target_state_update_decision[0].item()) else template[-2]
+                dynamic_mask = soft_token_template_mask[-1] if bool(target_state_update_decision[0].item()) else soft_token_template_mask[-2]
+                template_temporal = torch.cat([template[0], dynamic_template], dim=1)
+                soft_token_template_mask_temporal = torch.cat([soft_token_template_mask[0], dynamic_mask], dim=1)
+            else:
+                template_temporal = torch.cat(template[:2], dim=1)
+                soft_token_template_mask_temporal = torch.cat(soft_token_template_mask[:2], dim=1)
+
+        # x, aux_dict = self.backbone(z=template, x=search,
+        #                             soft_token_template_mask = soft_token_template_mask )
+        cls_prompts_pos = self.cls_prompts_pos.weight.unsqueeze(0)
+        x_pos_0 = torch.cat([cls_prompts_pos, self.backbone.pos_embed_z, self.backbone.pos_embed_x], dim=1)
+        # pos_embed = x_pos.transpose(0, 1).repeat(1, b0, 1)
+        x_pos = x_pos_0.repeat(b0*num_search, 1, 1)
+        x, aux_dict = self.forward_backbone(template_temporal, search, None, soft_token_template_mask_temporal,
+                                                 x_pos)
+        # forward Language branch
+        if training:
+            if exp_str:
+                text_features, text_subject_features, subject_infor_mask_pred, subject_infor_mask_gt  = self.forward_text(
+                    exp_str, num_search, exp_subject_mask, device=search.device)  # text_subject_features, subject_infor_mask_pred, subject_infor_mask_gt
+        else:
+            text_features = exp_str
+            text_subject_features = exp_subject_mask
+            subject_infor_mask_pred = None
+            subject_infor_mask_gt = None
+        if z_target is not None and z_target.shape[0] == b0 and num_search > 1:
+            z_target = torch.cat([z_target for _ in range(num_search)], dim=0)
+
+        batch_size = text_features.tensors.shape[0]
+        text_pos = self.text_pos(text_features) # [batch_size, length, c]
+        text_pos_0 = text_pos[:b0]
+        x_s_pos_item = x_pos_0.repeat(b0, 1, 1)[:, -self.feat_len_s:]
+        pre_temporal_pos = self.dy_template_pos_embed.weight.unsqueeze(1)
+        pre_temporal_pos = pre_temporal_pos.repeat(b0, 1, self.query_len)
+        pre_temporal_pos = pre_temporal_pos.view(b0, self.temporal_len * self.query_len, self.dim).contiguous()
+
+        # Forward temporal
+        xt_data = []
+        for temporal_index in range(num_search):
+            x_item = x[temporal_index * b0:(temporal_index + 1) * b0]
+
+            visual_prompts_token = x_item[:, :self.query_len, :]
+
+            ## heatmap by backbone feat
+            ## by attn
+            # attn_xz = attn[:, :, :-self.feat_len_s, -self.feat_len_s:]  #  b,h,l,l
+            # attn_xz_1 = attn_xz.mean(1).mean(1)
+            # # attn_xz = attn_xz.view(16, 16)
+            # # attn_weights_debug = attn_xz.detach().cpu().numpy()
+            x_f = x_item[:, -256:]
+            x_f1 = torch.matmul(x_f, x_f.permute(0, 2, 1).contiguous())
+            x_f = torch.matmul(x_f1, x_f)
+
+            z_f = x_item[:, :-256]
+
+            x_z = torch.matmul(x_f, z_f.permute(0, 2, 1).contiguous())
+            att_map = x_z.mean(-1)
+
+            tensor_min = torch.min(att_map)
+            tensor_max = torch.max(att_map)
+            # normalized_tensor = (s_vl_1 - tensor_min) / (tensor_max - tensor_min)
+            normalized_tensor = (tensor_max - att_map) / (tensor_max - tensor_min)
+
+            attn_xz = normalized_tensor.view(-1, 256,1).contiguous()
+
+            ### initialize & update memory
+            if training:
+                if temporal_index == 0:
+                    temporal_infor = []
+                    for _ in range(self.temporal_len):
+                        temporal_infor.append(visual_prompts_token)
+            else:
+                if first_frame_flag:
+                    temporal_infor = []
+                    for _ in range(self.temporal_len):
+                        temporal_infor.append(visual_prompts_token)
+
+            temporal_infor_data = torch.cat(temporal_infor, dim=1)
+
+            #### vl fusion  ############
+            ## L adjust
+            l_item_initial = text_features.tensors[temporal_index * b0:(temporal_index + 1) * b0]
+            l_item_subject = text_subject_features.tensors[temporal_index * b0:(temporal_index + 1) * b0]
+            l_mask_item_0 = text_features.mask[temporal_index * b0:(temporal_index + 1) * b0]
+            temporal_mask = torch.ones((l_mask_item_0.shape[0],self.temporal_len)).bool().to(l_mask_item_0.device)
+            l_mask_item = torch.cat([l_mask_item_0, temporal_mask],dim=1)
+
+            l_subject_temporal = torch.cat([l_item_subject,temporal_infor_data],dim=1)
+            l_subject_temporal_pos = torch.cat([text_pos_0,pre_temporal_pos ],dim=1)
+
+            l_item_update,_ = self.language_adjust([l_item_initial,l_subject_temporal],None,
+                                          text_pos_0,l_subject_temporal_pos,l_mask_item)
+            l_all = torch.cat([ l_item_initial,l_item_update ],dim=1)
+            x_s_item = x_item[:, -self.feat_len_s:]
+            x_s_item = self.vl_fusion(x_s_item,
+                                 l_all,
+                                 query_pos=x_pos_0[:, -self.feat_len_s:],
+                                 memory_pos=torch.cat([text_pos_0,text_pos_0],dim=1),
+                                 memory_key_padding_mask=torch.cat([l_mask_item_0,l_mask_item_0],dim=1),
+                                 need_weights=False)
+
+
+            #### cross_attention with temporal_infor
+            temporal_infor_update = self.visual_temporal_fusion(temporal_infor_data, x_s_item, attn_xz,pre_temporal_pos ,kv_pos= x_s_pos_item )
+            temporal_item = temporal_infor_update[:,-1,:].unsqueeze(1)
+
+            # STM
+            enc_opt = x_s_item
+            dec_opt = temporal_item.transpose(1, 2)
+            att = torch.matmul(enc_opt, dec_opt)
+            opt = (enc_opt.unsqueeze(-1) * att.unsqueeze(-2)).permute((0, 3, 2, 1)).contiguous()
+            bs, Nq, C, HW = opt.size()
+            opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s)
+            if z_target is not None:
+                z_item = z_target[temporal_index * b0:(temporal_index + 1) * b0]
+                opt_feat = self.target_state_encoder.modulate_feature(opt_feat, z_item)
+
+            xt_data.append(opt_feat)
+
+            ### update temporal infor
+            if training:
+                if temporal_index == 0:
+                    temporal_infor = []
+                    for _ in range(self.temporal_len):
+                        temporal_infor.append(temporal_item)
+                else:
+                    temporal_infor[:-1] = temporal_infor[1:]
+                    temporal_infor[-1] = temporal_item
+            else:
+                if first_frame_flag:
+                    temporal_infor = []
+                    for _ in range(self.temporal_len):
+                        temporal_infor.append(temporal_item)
+
+                else:
+                    temporal_infor[:-1] = temporal_infor[1:]
+                    temporal_infor[-1] = temporal_item
+
+
+        # Forward head
+        xt_data = torch.cat(xt_data,dim=0)
+        out = self.forward_head(xt_data, None)
+
+        out.update(aux_dict)
+        out['backbone_feat'] = x
+        out['subject_infor_mask_pred'] = subject_infor_mask_pred
+        out['subject_infor_mask_gt'] = subject_infor_mask_gt
+        out['target_state_update_decision'] = target_state_update_decision
+        out['qwen_format_loss'] = qwen_format_loss
+        out['qwen_teacher_loss'] = qwen_teacher_loss
+        out['qwen_teacher_labels'] = qwen_teacher_labels
+        out['qwen_teacher_outputs'] = qwen_teacher_outputs
+
+        if training == False:
+            out["temporal_infor"] = temporal_infor
+
+        return out
+
+    def forward_head(self, opt_feat, gt_score_map=None):
+        """
+        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
+        """
+
+        # enc_opt = cat_feature #[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)
+        # opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
+        # bs, Nq, C, HW = opt.size()
+        # opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s).contiguous()
+
+        bs = opt_feat.shape[0]
+        Nq = 1
+        # Head
+        if self.head_type == "CORNER":
+            # run the corner head
+            pred_box, score_map = self.box_head(opt_feat, True)
+            outputs_coord = box_xyxy_to_cxcywh(pred_box)
+            outputs_coord_new = outputs_coord.view(bs, Nq, 4).contiguous()
+            out = {'pred_boxes': outputs_coord_new,
+                   'score_map': score_map,
+                   }
+            return out
+
+        elif self.head_type == "CENTER":
+            # run the center head
+            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat, gt_score_map)
+            # outputs_coord = box_xyxy_to_cxcywh(bbox)
+
+            score_map = torch.cat([score_map_ctr, size_map, offset_map], dim=1)
+            confidence_pred = self.confidence_pred(score_map)
+
+            outputs_coord = bbox
+            outputs_coord_new = outputs_coord.view(bs, Nq, 4).contiguous()
+            out = {'pred_boxes': outputs_coord_new,
+                   'score_map': score_map_ctr,
+                   'size_map': size_map,
+                   'offset_map': offset_map,
+                   "confidence_pred": confidence_pred}
+            return out
+        else:
+            raise NotImplementedError
+
+    def forward_text(self, captions, num_search, exp_subject_mask, device):
+        tokenized = self.tokenizer(captions, padding=True, return_tensors="pt").to(device)
+        encoded_text = self.text_encoder(**tokenized)
+
+        text_attention_mask = tokenized.attention_mask.ne(1).bool()
+        # text_attention_mask: [batch_size, length]
+
+        text_features = encoded_text.last_hidden_state
+        text_features = self.text_adj(text_features)
+
+        encodings_infor = tokenized.encodings
+
+        subject_infor_mask_gt = None
+        if exp_subject_mask is not None:
+            # train: given the exp_subject_mask, used for generating  sub_index_gt
+            subject_infor_mask_gt = torch.zeros(text_attention_mask.shape[0], text_attention_mask.shape[1]).to(
+                text_features.device)
+
+            for item_index, item in enumerate(encodings_infor):
+                word_ids_item = item.word_ids
+                exp_subject_mask_item = exp_subject_mask[item_index]
+                text_index_list = []
+                for word_index, word_item in enumerate(word_ids_item):
+                    if word_item in exp_subject_mask_item:
+                        text_index_list.append(word_index)
+
+                subject_infor_mask_gt[item_index, text_index_list] = 1
+
+        subject_infor_mask_pred = self.text_sub_idnex_classifier(text_features)
+        subject_infor_mask_pred_1 = subject_infor_mask_pred.expand_as(text_features)
+
+        subject_infor = text_features * subject_infor_mask_pred_1
+
+        # (B,L,D) to (T,B,L,D)
+        text_features_t = []
+        text_attention_mask_t = []
+        text_subject_infor_t = []
+        for i in range(num_search):
+            text_features_t.append(text_features)
+            text_attention_mask_t.append(text_attention_mask)
+            text_subject_infor_t.append(subject_infor)
+
+        text_features = torch.cat(text_features_t, dim=0)
+        text_attention_mask = torch.cat(text_attention_mask_t, dim=0)
+        text_features = NestedTensor(text_features, text_attention_mask)
+        subject_infor = torch.cat(text_subject_infor_t, dim=0)
+        subject_infor = NestedTensor(subject_infor, text_attention_mask)
+
+        return text_features, subject_infor, subject_infor_mask_pred, subject_infor_mask_gt
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+def build_atctrack(cfg, training=True):
+    current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
+    pretrained_path = os.path.join(current_dir, '../../../resource/pretrained_models')
+
+    if cfg.MODEL.PRETRAIN_FILE  and training and ("ATCTrack" not in cfg.MODEL.PRETRAIN_FILE) :
+        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
+    else:
+        pretrained = ''
+
+
+    if cfg.MODEL.BACKBONE.TYPE == 'hivit_base_adaptor':
+        backbone = hivit_base(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+
+    elif cfg.MODEL.BACKBONE.TYPE == 'itpn_base':  # by this
+        backbone = fast_itpn_base_3324_patch16_224(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+    elif cfg.MODEL.BACKBONE.TYPE == 'itpn_large':  # by this
+        backbone = fast_itpn_large_2240_patch16_256(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+
+    else:
+        raise NotImplementedError
+
+    backbone.finetune_track(cfg=cfg,dim=hidden_dim, patch_start_index=patch_start_index)
+
+    box_head = build_box_head(cfg, hidden_dim)
+
+    # Build Text Encoder
+    roberta_path = _resolve_project_path(os.environ.get("ROBERTA_MODEL_PATH", os.path.join(pretrained_path, 'roberta-base')))
+    tokenizer = RobertaTokenizerFast.from_pretrained(roberta_path)  # load pretrained RoBERTa Tokenizer
+    text_encoder = RobertaModel.from_pretrained(roberta_path)  # load pretrained RoBERTa model
+
+
+    model = ATCTrack(
+        backbone,
+        box_head,
+        tokenizer,
+        text_encoder,
+        aux_loss=False,
+        head_type=cfg.MODEL.HEAD.TYPE,
+        dim = hidden_dim,
+        cfg=cfg
+    )
+
+    pretrained_checkpoint = _resolve_project_path(cfg.MODEL.PRETRAINED_PATH)
+    if  ("ATCTrack" in pretrained_checkpoint) and training:
+        checkpoint = torch.load(pretrained_checkpoint, map_location="cpu", weights_only=False)
+        ckpt = checkpoint["net"]
+        model_weight = {}
+        for k, v in ckpt.items():
+            model_weight[k] = v
+
+        missing_keys, unexpected_keys = model.load_state_dict(model_weight, strict=False)
+        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)
+
+
+    return model
+
+def load_pretrained(model, pretrained_path, strict=False):
+
+    model_ckpt = torch.load(pretrained_path, map_location="cpu")
+    state_dict = model_ckpt['net']
+    pos_st = state_dict['encoder.body.pos_embed']
+    pos_s = pos_st[:,:(pos_st.size(1) // 2)]
+    pos_t = pos_st[:,(pos_st.size(1) // 2):]
+    state_dict['encoder.body.pos_embed_search'] = pos_s
+    state_dict['encoder.body.pos_embed_template'] = pos_t
+    state_dict['encoder.body.patch_embed_interface.proj.weight'] = state_dict['encoder.body.patch_embed.proj.weight']
+    state_dict['encoder.body.patch_embed_interface.proj.bias'] = state_dict['encoder.body.patch_embed.proj.bias']
+    state_dict['decoder.embedding.prompt_embeddings.weight'] = model.state_dict()['decoder.embedding.prompt_embeddings.weight']
+    state_dict['decoder.embedding.prompt_embeddings.weight'][:] = state_dict['decoder.embedding.word_embeddings.weight'][-1]
+    del state_dict['encoder.body.pos_embed']
+    model.load_state_dict(state_dict, strict=strict)

lib/models/atctrack/decoder.py

@@ -0,0 +1,331 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+"""
+
+import copy
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+import torch.nn as nn
+
+
+class DecoderEmbeddings(nn.Module):
+    def __init__(self, vocab_size, instruct_vocab_size, hidden_dim, max_position_embeddings, dropout):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.instruct_vocab_size = instruct_vocab_size
+        self.hidden_dim = hidden_dim
+        self.word_embeddings = nn.Embedding(
+            vocab_size, hidden_dim)
+        self.prompt_embeddings = nn.Embedding(
+            instruct_vocab_size, hidden_dim)
+        self.position_embeddings = nn.Embedding(
+            max_position_embeddings, hidden_dim
+        )
+
+        self.LayerNorm = torch.nn.LayerNorm(
+            hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # input_embeds = self.word_embeddings(x)
+        # embeddings = input_embeds
+
+        use_word_embeddings = (x < self.vocab_size)
+        use_prompt_embeddings = ~use_word_embeddings
+
+        embeddings = torch.zeros([x.size(0),x.size(1),self.hidden_dim], dtype=torch.float32).to(x.device)
+        embeddings[use_word_embeddings] = self.word_embeddings(x[use_word_embeddings])
+        embeddings[use_prompt_embeddings] = self.prompt_embeddings(x[use_prompt_embeddings]-self.vocab_size)
+
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+
+        return embeddings
+
+
+class SeqTrackDecoder(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False, bins=1000, num_frames=9,
+                 instruct=True):
+        super().__init__()
+        self.bins = bins
+
+        self.instruct = instruct
+
+        self.instruct_tokens = {
+            'end': bins,
+            'lasot': bins+1,
+            'trackingnet': bins+1,
+            'got10k': bins+1,
+            'coco': bins+1,
+            'depthtrack': bins+2,
+            'lasher': bins+3,
+            'visevent': bins+4,
+            'otb99_lang': bins+5,
+            'refcocog': bins+5,
+            'tnl2k': bins+5,
+            'lasot_lang': bins+5
+        }
+        instruct_vocab_size = 4 # should be consistent with new tokens in self.instruct_tokens
+
+        self.num_frames = num_frames
+        self.num_coordinates = 4 # [x,y,w,h]
+        max_position_embeddings = (self.num_coordinates+1) * num_frames
+        self.embedding = DecoderEmbeddings(bins+2, instruct_vocab_size, d_model, max_position_embeddings, dropout)
+
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        decoder_norm = nn.LayerNorm(d_model)
+        self.body = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
+                                          return_intermediate=return_intermediate_dec)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, pos_embed, seq):
+        # flatten NxCxHxW to HWxNxC
+        n, bs, c = src.shape
+        if not self.instruct:
+            seq[:, 0] = self.bins+1
+        tgt = self.embedding(seq).permute(1, 0, 2)
+
+        query_embed = self.embedding.position_embeddings.weight.unsqueeze(1)
+        query_embed = query_embed.repeat(1, bs, 1)
+
+        memory = src
+
+        tgt_mask = generate_square_subsequent_mask(len(tgt)).to(tgt.device) #generate the causal mask
+
+        hs = self.body(tgt, memory, pos=pos_embed, query_pos=query_embed[:len(tgt)],
+                       tgt_mask=tgt_mask, memory_mask=None)
+
+        return hs.transpose(1, 2)
+
+
+    def inference(self, src, pos_embed, seq, vocab_embed,
+                  window, seq_format):
+        if not self.instruct:
+            seq[:, 0] = self.bins+1
+        # flatten NxCxHxW to HWxNxC
+        n, bs, c = src.shape
+        memory = src
+        confidence_list = []
+        box_pos = [0, 1, 2, 3] # the position of bounding box
+        center_pos = [0, 1]  # the position of x_center and y_center
+        if seq_format == 'whxy':
+            center_pos = [2, 3]
+
+        for i in range(self.num_coordinates): # only cycle 4 times, because we do not need to predict the end token during inference
+            tgt = self.embedding(seq).permute(1, 0, 2)
+            query_embed = self.embedding.position_embeddings.weight.unsqueeze(1)
+            query_embed = query_embed.repeat(1, bs, 1)
+            tgt_mask = generate_square_subsequent_mask(len(tgt)).to(tgt.device)
+
+            hs = self.body(tgt, memory, pos=pos_embed[:len(memory)], query_pos=query_embed[:len(tgt)],
+                              tgt_mask=tgt_mask, memory_mask=None)
+
+            # embedding --> likelihood
+            out = vocab_embed(hs.transpose(1, 2)[-1, :, -1, :])
+            out = out.softmax(-1)
+
+            if i in box_pos:
+                out = out[:, :self.bins] # only include the coordinate values' confidence
+
+            if ((i in center_pos) and (window!=None)):
+                out = out * window # window penalty
+
+            confidence, token_generated = out.topk(dim=-1, k=1)
+            seq = torch.cat([seq, token_generated], dim=-1)
+            confidence_list.append(confidence)
+
+        out_dict = {}
+        out_dict['pred_boxes'] = seq[:, -self.num_coordinates:] # Discard the START token, only get the bounding box
+        out_dict['confidence'] = torch.cat(confidence_list, dim=-1)[:, :]
+
+        return out_dict
+
+
+
+
+def generate_square_subsequent_mask(sz):
+    r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+        Unmasked positions are filled with float(0.0).
+    """
+
+    #each token only can see tokens before them
+    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+    mask = mask.float().masked_fill(mask == 0, float(
+        '-inf')).masked_fill(mask == 1, float(0.0))
+    return mask
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+
+        intermediate = []
+
+        for layer in self.layers:
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos)
+
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output.unsqueeze(0)
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(q, k, tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(self.with_pos_embed(tgt, query_pos),
+                                   self.with_pos_embed(memory, pos),
+                                   memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+
+        return tgt
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(self.with_pos_embed(tgt2, query_pos),
+                                   self.with_pos_embed(memory, pos),
+                                   memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_decoder(cfg):
+    return SeqTrackDecoder(
+        d_model=cfg.MODEL.HIDDEN_DIM,
+        dropout=cfg.MODEL.DECODER.DROPOUT,
+        nhead=cfg.MODEL.DECODER.NHEADS,
+        dim_feedforward=cfg.MODEL.DECODER.DIM_FEEDFORWARD,
+        num_decoder_layers=cfg.MODEL.DECODER.DEC_LAYERS,
+        normalize_before=cfg.MODEL.DECODER.PRE_NORM,
+        return_intermediate_dec=False,
+        bins=cfg.MODEL.BINS,
+        num_frames=cfg.DATA.SEARCH.NUMBER,
+        instruct=cfg.MODEL.DECODER.INSTRUCT
+    )
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

lib/models/atctrack/encoder.py

@@ -0,0 +1,100 @@
+"""
+Encoder modules: we use ViT for the encoder.
+"""
+
+from torch import nn
+from lib.utils.misc import is_main_process
+from lib.models.atctrack import vit as vit_module
+from lib.models.atctrack import vit_mm as vitmm_module
+
+
+
+
+class EncoderBase(nn.Module):
+
+    def __init__(self, encoder: nn.Module, train_encoder: bool, open_layers: list, num_channels: int):
+        super().__init__()
+        open_blocks = open_layers[2:]
+        open_items = open_layers[0:2]
+        for name, parameter in encoder.named_parameters():
+
+            if not train_encoder:
+                freeze = True
+                for open_block in open_blocks:
+                    if open_block in name:
+                        freeze = False
+                if name in open_items:
+                    freeze = False
+                if freeze == True:
+                    parameter.requires_grad_(False)  # here should allow users to specify which layers to freeze !
+
+        self.body = encoder
+        self.num_channels = num_channels
+
+    def forward(self, template_list, search_list, text_src, seq):
+        xs = self.body(template_list, search_list, text_src, seq)
+        return xs
+
+    def forward_rgb(self, template_list, search_list):
+        xs = self.body.forward_rgb(template_list, search_list)
+        return xs
+
+
+class Encoder(EncoderBase):
+    """ViT encoder."""
+    def __init__(self, name: str,
+                 train_encoder: bool,
+                 pretrain_type: str,
+                 search_size: int,
+                 search_number: int,
+                 template_size: int,
+                 template_number: int,
+                 open_layers: list,
+                 cfg=None):
+        if "vitmm" in name.lower():
+            encoder = getattr(vitmm_module, name)(pretrained=is_main_process(), pretrain_type=pretrain_type,
+                                                  search_size=search_size, template_size=template_size,
+                                                  search_number=search_number, template_number=template_number,
+                                                  drop_path_rate=cfg.MODEL.ENCODER.DROP_PATH,
+                                                  use_checkpoint=cfg.MODEL.ENCODER.USE_CHECKPOINT,
+                                                  interface_type=cfg.MODEL.INTERFACE_TYPE,
+                                                  interface_dim=cfg.MODEL.INTERFACE_DIM,
+                                                  instruct=cfg.MODEL.ENCODER.INSTRUCT)
+            if "_base_" in name:
+                num_channels = 768
+            elif "_large_" in name:
+                num_channels = 1024
+            elif "_huge_" in name:
+                num_channels = 1280
+            else:
+                num_channels = 768
+        elif "vit" in name.lower():
+            encoder = getattr(vit_module, name)(pretrained=is_main_process(), pretrain_type=pretrain_type,
+                                                       search_size=search_size, template_size=template_size,
+                                                       search_number=search_number, template_number=template_number,
+                                                       drop_path_rate=cfg.MODEL.ENCODER.DROP_PATH,
+                                                       use_checkpoint=cfg.MODEL.ENCODER.USE_CHECKPOINT
+                                                      )
+            if "_base_" in name:
+                num_channels = 768
+            elif "_large_" in name:
+                num_channels = 1024
+            elif "_huge_" in name:
+                num_channels = 1280
+            else:
+                num_channels = 768
+
+        else:
+            raise ValueError()
+        super().__init__(encoder, train_encoder, open_layers, num_channels)
+
+
+
+def build_encoder(cfg):
+    train_encoder = (cfg.TRAIN.ENCODER_MULTIPLIER > 0) and (cfg.TRAIN.FREEZE_ENCODER == False)
+    encoder = Encoder(cfg.MODEL.ENCODER.TYPE, train_encoder,
+                      cfg.MODEL.ENCODER.PRETRAIN_TYPE,
+                      cfg.DATA.SEARCH.SIZE, cfg.DATA.SEARCH.NUMBER,
+                      cfg.DATA.TEMPLATE.SIZE, cfg.DATA.TEMPLATE.NUMBER,
+                      cfg.TRAIN.ENCODER_OPEN, cfg)
+    return encoder

lib/models/atctrack/language_model.py

@@ -0,0 +1,320 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Backbone modules.
+"""
+import os.path
+import shutil
+import tarfile
+import tempfile
+
+import torch
+from pytorch_pretrained_bert import BertConfig, cached_path, CONFIG_NAME, WEIGHTS_NAME, load_tf_weights_in_bert
+from pytorch_pretrained_bert.modeling import BertLayerNorm, PRETRAINED_MODEL_ARCHIVE_MAP, logger, BERT_CONFIG_NAME, \
+    BertEmbeddings, BertEncoder
+from pytorch_pretrained_bert.modeling_transfo_xl import TF_WEIGHTS_NAME
+from torch import nn
+
+from lib.utils.misc import NestedTensor
+
+
+class BertPreTrainedModel(nn.Module):
+    """ An abstract class to handle weights initialization and
+        a simple interface for dowloading and loading pretrained checkpoints.
+    """
+    def __init__(self, config, *inputs, **kwargs):
+        super(BertPreTrainedModel, self).__init__()
+        if not isinstance(config, BertConfig):
+            raise ValueError(
+                "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
+                "To create a model from a Google pretrained model use "
+                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    self.__class__.__name__, self.__class__.__name__
+                ))
+        self.config = config
+
+    def init_bert_weights(self, module):
+        """ Initialize the weights.
+        """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, BertLayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+        """
+        Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.
+        Download and cache the pre-trained model file if needed.
+
+        Params:
+            pretrained_model_name_or_path: either:
+                - a str with the name of a pre-trained model to load selected in the list of:
+                    . `bert-base-uncased`
+                    . `bert-large-uncased`
+                    . `bert-base-cased`
+                    . `bert-large-cased`
+                    . `bert-base-multilingual-uncased`
+                    . `bert-base-multilingual-cased`
+                    . `bert-base-chinese`
+                - a path or url to a pretrained model archive containing:
+                    . `bert_config.json` a configuration file for the model
+                    . `pytorch_model.bin` a PyTorch dump of a BertForPreTraining instance
+                - a path or url to a pretrained model archive containing:
+                    . `bert_config.json` a configuration file for the model
+                    . `model.chkpt` a TensorFlow checkpoint
+            from_tf: should we load the weights from a locally saved TensorFlow checkpoint
+            cache_dir: an optional path to a folder in which the pre-trained checkpoints will be cached.
+            state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of Google pre-trained checkpoints
+            *inputs, **kwargs: additional input for the specific Bert class
+                (ex: num_labels for BertForSequenceClassification)
+        """
+        state_dict = kwargs.get('state_dict', None)
+        kwargs.pop('state_dict', None)
+        cache_dir = kwargs.get('cache_dir', None)
+        kwargs.pop('cache_dir', None)
+        from_tf = kwargs.get('from_tf', False)
+        kwargs.pop('from_tf', None)
+
+        if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
+            archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
+        else:
+            archive_file = pretrained_model_name_or_path
+        # redirect to the cache, if necessary
+        try:
+            resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
+        except EnvironmentError:
+            logger.error(
+                "Model name '{}' was not found in model name list ({}). "
+                "We assumed '{}' was a path or url but couldn't find any file "
+                "associated to this path or url.".format(
+                    pretrained_model_name_or_path,
+                    ', '.join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()),
+                    archive_file))
+            return None
+        if resolved_archive_file == archive_file:
+            logger.info("loading archive file {}".format(archive_file))
+        else:
+            logger.info("loading archive file {} from cache at {}".format(
+                archive_file, resolved_archive_file))
+        tempdir = None
+        if os.path.isdir(resolved_archive_file) or from_tf:
+            serialization_dir = resolved_archive_file
+        else:
+            # Extract archive to temp dir
+            tempdir = tempfile.mkdtemp()
+            logger.info("extracting archive file {} to temp dir {}".format(
+                resolved_archive_file, tempdir))
+            with tarfile.open(resolved_archive_file, 'r:gz') as archive:
+                archive.extractall(tempdir)
+            serialization_dir = tempdir
+        # Load config
+        config_file = os.path.join(serialization_dir, CONFIG_NAME)
+        if not os.path.exists(config_file):
+            # Backward compatibility with old naming format
+            config_file = os.path.join(serialization_dir, BERT_CONFIG_NAME)
+        config = BertConfig.from_json_file(config_file)
+        logger.info("Model config {}".format(config))
+        # Instantiate model.
+        model = cls(config, *inputs, **kwargs)
+        if state_dict is None and not from_tf:
+            weights_path = os.path.join(serialization_dir, WEIGHTS_NAME)
+            state_dict = torch.load(weights_path, map_location='cpu')
+        if tempdir:
+            # Clean up temp dir
+            shutil.rmtree(tempdir)
+        if from_tf:
+            # Directly load from a TensorFlow checkpoint
+            weights_path = os.path.join(serialization_dir, TF_WEIGHTS_NAME)
+            return load_tf_weights_in_bert(model, weights_path)
+        # Load from a PyTorch state_dict
+        old_keys = []
+        new_keys = []
+        for key in state_dict.keys():
+            new_key = None
+            if 'gamma' in key:
+                new_key = key.replace('gamma', 'weight')
+            if 'beta' in key:
+                new_key = key.replace('beta', 'bias')
+            if new_key:
+                old_keys.append(key)
+                new_keys.append(new_key)
+        for old_key, new_key in zip(old_keys, new_keys):
+            state_dict[new_key] = state_dict.pop(old_key)
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+        # copy state_dict so _load_from_state_dict can modify it
+        metadata = getattr(state_dict, '_metadata', None)
+        state_dict = state_dict.copy()
+        if metadata is not None:
+            state_dict._metadata = metadata
+
+        def load(module, prefix=''):
+            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+            module._load_from_state_dict(
+                state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+            for name, child in module._modules.items():
+                if child is not None:
+                    load(child, prefix + name + '.')
+        start_prefix = ''
+        if not hasattr(model, 'bert') and any(s.startswith('bert.') for s in state_dict.keys()):
+            start_prefix = 'bert.'
+        load(model, prefix=start_prefix)
+        if len(missing_keys) > 0:
+            logger.info("Weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, missing_keys))
+        if len(unexpected_keys) > 0:
+            logger.info("Weights from pretrained model not used in {}: {}".format(
+                model.__class__.__name__, unexpected_keys))
+        if len(error_msgs) > 0:
+            raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
+                               model.__class__.__name__, "\n\t".join(error_msgs)))
+        return model
+
+
+class BertModel(BertPreTrainedModel):
+    """BERT model ("Bidirectional Embedding Representations from a Transformer").
+
+    Params:
+        config: a BertConfig class instance with the configuration to build a new model
+
+    Inputs:
+        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
+            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
+            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
+        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
+            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
+            a `sentence B` token (see BERT paper for more details).
+        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
+            input sequence length in the current batch. It's the mask that we typically use for attention when
+            a batch has varying length sentences.
+        `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
+
+    Outputs: Tuple of (encoded_layers, pooled_output)
+        `encoded_layers`: controled by `output_all_encoded_layers` argument:
+            - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
+                of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
+                encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
+            - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
+                to the last attention block of shape [batch_size, sequence_length, hidden_size],
+        `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
+            classifier pretrained on top of the hidden state associated to the first character of the
+            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
+
+    Example usage:
+    ```python
+    # Already been converted into WordPiece token ids
+    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
+
+    config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
+        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
+
+    model = modeling.BertModel(config=config)
+    all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
+    ```
+    """
+    def __init__(self, config):
+        super(BertModel, self).__init__(config)
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        # self.pooler = BertPooler(config)
+        self.apply(self.init_bert_weights)
+
+    def forward(self, input_ids, token_type_ids=None, attention_mask=None, output_all_encoded_layers=True):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        # We create a 3D attention mask from a 2D tensor mask.
+        # Sizes are [batch_size, 1, 1, to_seq_length]
+        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+        # this attention mask is more simple than the triangular masking of causal attention
+        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        embedding_output = self.embeddings(input_ids, token_type_ids)
+        encoded_layers = self.encoder(embedding_output,
+                                      extended_attention_mask,
+                                      output_all_encoded_layers=output_all_encoded_layers)
+        # sequence_output = encoded_layers[-1]
+        # pooled_output = self.pooler(sequence_output)
+        if not output_all_encoded_layers:
+            encoded_layers = encoded_layers[-1]
+        return encoded_layers
+
+
+class BERT(nn.Module):
+    def __init__(self, name: str, path: str, train_bert: bool, hidden_dim: int, max_len: int, enc_num):
+        super().__init__()
+        if name == 'bert-base-uncased':
+            self.num_channels = 768
+        else:
+            self.num_channels = 1024
+        self.enc_num = enc_num
+        if path is not None and os.path.exists(path):
+            self.bert = BertModel.from_pretrained(path)
+        else:
+            self.bert = BertModel.from_pretrained(name)
+
+        if not train_bert:
+            print('Language Model Bert has been frozen!')
+            for parameter in self.bert.parameters():
+                parameter.requires_grad_(False)
+        # print(self.bert)
+
+    def forward(self, tensor_list: NestedTensor):
+
+        if self.enc_num > 0:
+            all_encoder_layers = self.bert(tensor_list.tensors, token_type_ids=None, attention_mask=tensor_list.mask)
+            # use the output of the X-th transformer encoder layers
+            xs = all_encoder_layers[self.enc_num - 1]
+        else:
+            xs = self.bert.embeddings.word_embeddings(tensor_list.tensors)
+
+        mask = tensor_list.mask.to(torch.bool)
+        mask = ~mask
+        out = NestedTensor(xs, mask)
+
+        return out
+
+
+# def build_bert(cfg):
+#     # position_embedding = build_position_encoding(cfg)
+#     train_bert = cfg.MODEL.LANGUAGE.BERT.LR > 0
+#     bert = BERT(cfg.MODEL.LANGUAGE.TYPE, cfg.MODEL.LANGUAGE.PATH, train_bert, cfg.MODEL.LANGUAGE.BERT.HIDDEN_DIM,
+#                 cfg.MODEL.LANGUAGE.BERT.MAX_QUERY_LEN, cfg.MODEL.LANGUAGE.BERT.ENC_NUM)
+#     # model = Joiner(bert, position_embedding)
+#     # model.num_channels = bert.num_channels
+#     return bert
+
+def build_bert(cfg):
+    # position_embedding = build_position_encoding(cfg)
+    # train_bert = cfg.MODEL.LANGUAGE.BERT.LR > 0
+    train_bert = cfg.TRAIN.LR > 0
+    bert_type = cfg.MODEL.LANGUAGE.IMPLEMENT
+    if bert_type == "pytorch":
+        bert_model = BERT(cfg.MODEL.LANGUAGE.TYPE, cfg.MODEL.LANGUAGE.PATH, train_bert,
+                          cfg.MODEL.LANGUAGE.BERT.HIDDEN_DIM,
+                          cfg.MODEL.LANGUAGE.BERT.MAX_QUERY_LEN, cfg.MODEL.LANGUAGE.BERT.ENC_NUM)
+    else:
+        raise ValueError("Undefined BERT TYPE '%s'" % bert_type)
+    return bert_model
+

lib/models/atctrack/utils.py

@@ -0,0 +1,64 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+def combine_tokens(template_tokens, search_tokens, mode='direct', return_res=False):
+    # [B, HW, C]
+    len_t = template_tokens.shape[1]
+    len_s = search_tokens.shape[1]
+
+    if mode == 'direct':
+        merged_feature = torch.cat((template_tokens, search_tokens), dim=1)
+    elif mode == 'template_central':
+        central_pivot = len_s // 2
+        first_half = search_tokens[:, :central_pivot, :]
+        second_half = search_tokens[:, central_pivot:, :]
+        merged_feature = torch.cat((first_half, template_tokens, second_half), dim=1)
+    elif mode == 'partition':
+        feat_size_s = int(math.sqrt(len_s))
+        feat_size_t = int(math.sqrt(len_t))
+        window_size = math.ceil(feat_size_t / 2.)
+        # pad feature maps to multiples of window size
+        B, _, C = template_tokens.shape
+        H = W = feat_size_t
+        template_tokens = template_tokens.view(B, H, W, C)
+        pad_l = pad_b = pad_r = 0
+        # pad_r = (window_size - W % window_size) % window_size
+        pad_t = (window_size - H % window_size) % window_size
+        template_tokens = F.pad(template_tokens, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = template_tokens.shape
+        template_tokens = template_tokens.view(B, Hp // window_size, window_size, W, C)
+        template_tokens = torch.cat([template_tokens[:, 0, ...], template_tokens[:, 1, ...]], dim=2)
+        _, Hc, Wc, _ = template_tokens.shape
+        template_tokens = template_tokens.view(B, -1, C)
+        merged_feature = torch.cat([template_tokens, search_tokens], dim=1)
+
+        # calculate new h and w, which may be useful for SwinT or others
+        merged_h, merged_w = feat_size_s + Hc, feat_size_s
+        if return_res:
+            return merged_feature, merged_h, merged_w
+
+    else:
+        raise NotImplementedError
+
+    return merged_feature
+
+'''
+add token transfer to feature
+'''
+def token2feature(tokens):
+    B,L,D=tokens.shape
+    H=W=int(L**0.5)
+    x = tokens.permute(0, 2, 1).view(B, D, W, H).contiguous()
+    return x
+
+
+'''
+feature2token
+'''
+def feature2token(x):
+    B,C,W,H = x.shape
+    L = W*H
+    tokens = x.view(B, C, L).permute(0, 2, 1).contiguous()
+    return tokens

lib/models/atctrack/vit.py

@@ -0,0 +1,466 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Status/TODO:
+* Models updated to be compatible with official impl. Args added to support backward compat for old PyTorch weights.
+* Weights ported from official jax impl for 384x384 base and small models, 16x16 and 32x32 patches.
+* Trained (supervised on ImageNet-1k) my custom 'small' patch model to 77.9, 'base' to 79.4 top-1 with this code.
+* Hopefully find time and GPUs for SSL or unsupervised pretraining on OpenImages w/ ImageNet fine-tune in future.
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch.utils.checkpoint as checkpoint
+import torch.nn.functional as F
+import math
+from functools import partial
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = {
+    # patch models
+    # mae ViT-B/16-224 pre-trained model
+    'vit_base_patch16_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_base_patch16_224_default': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    # mae ViT-L/16-224 pre-trained model
+    'vit_large_patch16_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    # mae ViT-H/14-224 pre-trained model
+    'vit_huge_patch14_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_small_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth',
+    ),
+    'vit_base_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+    ),
+    'vit_base_patch16_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_base_patch32_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p32_384-830016f5.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_large_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+    'vit_large_patch16_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_384-b3be5167.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_large_patch32_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_huge_patch16_224': _cfg(),
+    'vit_huge_patch32_384': _cfg(input_size=(3, 384, 384)),
+    # hybrid models
+    'vit_small_resnet26d_224': _cfg(),
+    'vit_small_resnet50d_s3_224': _cfg(),
+    'vit_base_resnet26d_224': _cfg(),
+    'vit_base_resnet50d_224': _cfg(),
+}
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        self.img_size = img_size
+        self.backbone = backbone
+        if feature_size is None:
+            with torch.no_grad():
+                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
+                # map for all networks, the feature metadata has reliable channel and stride info, but using
+                # stride to calc feature dim requires info about padding of each stage that isn't captured.
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        self.num_patches = feature_size[0] * feature_size[1]
+        self.proj = nn.Linear(feature_dim, embed_dim)
+
+    def forward(self, x):
+        x = self.backbone(x)[-1]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, search_size=384, template_size=192,
+                 patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm,
+                 search_number=1, template_number=1, use_checkpoint=False):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.embed_dim_list = [embed_dim]
+        self.num_search = search_number
+        self.num_template = template_number
+
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+
+        self.num_patches_search = (search_size // patch_size) * (search_size // patch_size)
+        self.num_patches_template = (template_size // patch_size) * (template_size // patch_size)
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches_search + self.num_patches_template, embed_dim))
+        # self.pos_embed_search = nn.Parameter(torch.zeros(1, self.num_patches_search, embed_dim))
+        # self.pos_embed_template = nn.Parameter(torch.zeros(1, self.num_patches_template, embed_dim))
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+
+        self.norm = norm_layer(embed_dim)
+        trunc_normal_(self.pos_embed, std=.02)
+        # trunc_normal_(self.pos_embed_search, std=.02)
+        # trunc_normal_(self.pos_embed_template, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, template_list, search_list):
+        num_template = len(template_list)
+        num_search = len(search_list)
+
+        z = torch.stack(template_list, dim=1)#(b,n,c,h,w)
+        z = z.view(-1, *z.size()[2:])#(bn,c,h,w)
+        z = self.patch_embed(z)
+        z = z + self.pos_embed[:, self.num_patches_search:, :]
+        z = z.view(-1,num_template,z.size(-2),z.size(-1))
+        z = z.reshape(z.size(0),-1,z.size(-1))
+
+        x = torch.stack(search_list, dim=1)#(b,n,c,h,w)
+        x = x.view(-1, *x.size()[2:])#(bn,c,h,w)
+        x = self.patch_embed(x)
+        x = x + self.pos_embed[:, :self.num_patches_search, :]
+        x = x.view(-1,num_search,x.size(-2),x.size(-1))
+        x = x.reshape(x.size(0),-1,x.size(-1))
+
+        xz = torch.cat([x, z], dim=1)
+
+        xz = self.pos_drop(xz)
+
+        for blk in self.blocks:   #batch is the first dimension.
+            if self.use_checkpoint:
+                xz = checkpoint.checkpoint(blk, xz)
+            else:
+                xz = blk(xz)
+
+        xz = self.norm(xz) # B,N,C
+        return xz
+
+    def forward(self, template_list, search_list):
+        xz = self.forward_features(template_list, search_list)
+        out=[xz]
+        return out
+
+
+def _conv_filter(state_dict, patch_size=16):
+    """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    out_dict = {}
+    for k, v in state_dict.items():
+        if 'patch_embed.proj.weight' in k:
+            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+        out_dict[k] = v
+    return out_dict
+
+
+@register_model
+def vit_base_patch16(pretrained=False, pretrain_type='default',
+                         search_size=384, template_size=192, **kwargs):
+    patch_size = 16
+    model = VisionTransformer(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_base_patch16_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model, pretrain_type,
+                        num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+@register_model
+def vit_large_patch16(pretrained=False, pretrain_type='default',
+                              search_size=384, template_size=192, **kwargs):
+    patch_size = 16
+    model = VisionTransformer(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_large_patch16_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model, pretrain_type, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+@register_model
+def vit_huge_patch14(pretrained=False, pretrain_type='default',
+                             search_size=364, template_size=182, **kwargs):
+    patch_size = 14
+    model = VisionTransformer(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_huge_patch14_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model,
+                        pretrain_type, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+def load_pretrained(model, pretrain_type='default', cfg=None, num_classes=1000, in_chans=3, filter_fn=None, strict=True):
+    if cfg is None:
+        cfg = getattr(model, 'default_cfg')
+    if cfg is None or 'url' not in cfg or not cfg['url']:
+        print("Pretrained model URL is invalid, using random initialization.")
+        return
+
+    state_dict = model_zoo.load_url(cfg['url'], progress=False, map_location='cpu')
+    if pretrain_type == 'mae':
+        state_dict = state_dict['model']
+
+    if filter_fn is not None:
+        state_dict = filter_fn(state_dict)
+
+    if in_chans == 1:
+        conv1_name = cfg['first_conv']
+        print('Converting first conv (%s) pretrained weights from 3 to 1 channel' % conv1_name)
+        conv1_weight = state_dict[conv1_name + '.weight']
+        # Some weights are in torch.half, ensure it's float for sum on CPU
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I > 3:
+            assert conv1_weight.shape[1] % 3 == 0
+            # For models with space2depth stems
+            conv1_weight = conv1_weight.reshape(O, I // 3, 3, J, K)
+            conv1_weight = conv1_weight.sum(dim=2, keepdim=False)
+        else:
+            conv1_weight = conv1_weight.sum(dim=1, keepdim=True)
+        conv1_weight = conv1_weight.to(conv1_type)
+        state_dict[conv1_name + '.weight'] = conv1_weight
+    elif in_chans != 3:
+        conv1_name = cfg['first_conv']
+        conv1_weight = state_dict[conv1_name + '.weight']
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I != 3:
+            print('Deleting first conv (%s) from pretrained weights.' % conv1_name)
+            del state_dict[conv1_name + '.weight']
+            strict = False
+        else:
+            # NOTE this strategy should be better than random init, but there could be other combinations of
+            # the original RGB input layer weights that'd work better for specific cases.
+            print('Repeating first conv (%s) weights in channel dim.' % conv1_name)
+            repeat = int(math.ceil(in_chans / 3))
+            conv1_weight = conv1_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
+            conv1_weight *= (3 / float(in_chans))
+            conv1_weight = conv1_weight.to(conv1_type)
+            state_dict[conv1_name + '.weight'] = conv1_weight
+
+    classifier_name = cfg['classifier']
+    if pretrain_type == "mae":
+        pass
+    elif num_classes == 1000 and cfg['num_classes'] == 1001:
+        # special case for imagenet trained models with extra background class in pretrained weights
+        classifier_weight = state_dict[classifier_name + '.weight']
+        state_dict[classifier_name + '.weight'] = classifier_weight[1:]
+        classifier_bias = state_dict[classifier_name + '.bias']
+        state_dict[classifier_name + '.bias'] = classifier_bias[1:]
+    elif num_classes != cfg['num_classes']:
+        # completely discard fully connected for all other differences between pretrained and created model
+        del state_dict[classifier_name + '.weight']
+        del state_dict[classifier_name + '.bias']
+
+    # adjust position encoding
+    pe = state_dict['pos_embed'][:,1:,:]
+    b_pe, hw_pe, c_pe = pe.shape
+    side_pe = int(math.sqrt(hw_pe))
+    side_num_patches_search = int(math.sqrt(model.num_patches_search))
+    side_num_patches_template = int(math.sqrt(model.num_patches_template))
+    pe_2D = pe.reshape([b_pe, side_pe, side_pe, c_pe]).permute([0,3,1,2])  #b,c,h,w
+    if side_pe != side_num_patches_search:
+        pe_s_2D = nn.functional.interpolate(pe_2D, [side_num_patches_search, side_num_patches_search], align_corners=True, mode='bicubic')
+        pe_s = torch.flatten(pe_s_2D.permute([0,2,3,1]),1,2)
+    else:
+        pe_s = pe
+    if side_pe != side_num_patches_template:
+        pe_t_2D = nn.functional.interpolate(pe_2D, [side_num_patches_template, side_num_patches_template], align_corners=True, mode='bicubic')
+        pe_t = torch.flatten(pe_t_2D.permute([0, 2, 3, 1]), 1, 2)
+    else:
+        pe_t = pe
+    pe_xz = torch.cat((pe_s, pe_t), dim=1)
+    state_dict['pos_embed'] = pe_xz
+    # state_dict['pos_embed_template'] = pe_t
+    # state_dict['pos_embed_search'] = pe_s
+    # del state_dict['pos_embed']
+    del state_dict['cls_token']
+
+    model.load_state_dict(state_dict, strict=strict)
+

lib/models/atctrack/vit_mm.py

@@ -0,0 +1,575 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Status/TODO:
+* Models updated to be compatible with official impl. Args added to support backward compat for old PyTorch weights.
+* Weights ported from official jax impl for 384x384 base and small models, 16x16 and 32x32 patches.
+* Trained (supervised on ImageNet-1k) my custom 'small' patch model to 77.9, 'base' to 79.4 top-1 with this code.
+* Hopefully find time and GPUs for SSL or unsupervised pretraining on OpenImages w/ ImageNet fine-tune in future.
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch.utils.checkpoint as checkpoint
+import torch.nn.functional as F
+import math
+from functools import partial
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+from .utils import combine_tokens, token2feature, feature2token
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = {
+    # patch models
+    # mae ViT-B/16-224 pre-trained model
+    'vit_base_patch16_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_base_patch16_224_default': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    # mae ViT-L/16-224 pre-trained model
+    'vit_large_patch16_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    # mae ViT-H/14-224 pre-trained model
+    'vit_huge_patch14_224_mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth',
+        input_size=(3, 224, 224), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_small_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth',
+    ),
+    'vit_base_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+    ),
+    'vit_base_patch16_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_base_patch32_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p32_384-830016f5.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_large_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+    'vit_large_patch16_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_384-b3be5167.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_large_patch32_384': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pth',
+        input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),
+    'vit_huge_patch16_224': _cfg(),
+    'vit_huge_patch32_384': _cfg(input_size=(3, 384, 384)),
+    # hybrid models
+    'vit_small_resnet26d_224': _cfg(),
+    'vit_small_resnet50d_s3_224': _cfg(),
+    'vit_base_resnet26d_224': _cfg(),
+    'vit_base_resnet50d_224': _cfg(),
+}
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+class Interface_block(nn.Module, ):
+    def __init__(self, inplanes=None, hide_channel=None):
+        super(Interface_block, self).__init__()
+        self.conv0_0 = nn.Conv2d(in_channels=inplanes, out_channels=hide_channel, kernel_size=1, stride=1, padding=0)
+        self.conv0_1 = nn.Conv2d(in_channels=inplanes, out_channels=hide_channel, kernel_size=1, stride=1, padding=0)
+        self.conv1x1 = nn.Conv2d(in_channels=hide_channel, out_channels=inplanes, kernel_size=1, stride=1, padding=0)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x):
+        """ Forward pass with input x. """
+        B, C, W, H = x.shape
+        x0 = x[:, 0:int(C/2), :, :].contiguous()
+        x0 = self.conv0_0(x0)
+        x1 = x[:, int(C/2):, :, :].contiguous()
+        x1 = self.conv0_1(x1)
+        x0 = x0 + x1
+        return self.conv1x1(x0)
+
+class VisionTransformerMM(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, search_size=384, template_size=192,
+                 patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm,
+                 search_number=1, template_number=1, use_checkpoint=False,
+                 interface_type=None, interface_dim=8, instruct=True):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.embed_dim_list = [embed_dim]
+        self.num_search = search_number
+        self.num_template = template_number
+
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+
+        self.patch_embed_interface = PatchEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+
+        self.num_patches_search = (search_size // patch_size) * (search_size // patch_size)
+        self.num_patches_template = (template_size // patch_size) * (template_size // patch_size)
+
+        # self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches_search + self.num_patches_template, embed_dim))
+        self.pos_embed_search = nn.Parameter(torch.zeros(1, self.num_patches_search, embed_dim))
+        self.pos_embed_template = nn.Parameter(torch.zeros(1, self.num_patches_template, embed_dim))
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # for multi-modal
+        self.interface_type = interface_type
+        '''interface parameters'''
+        if self.interface_type in ['low-rank_add']:
+            interface_blocks = []
+            block_nums = depth
+            for i in range(block_nums):
+                if self.interface_type == 'low-rank_add':
+                    interface_blocks.append(Interface_block(inplanes=embed_dim, hide_channel=interface_dim))
+                else:
+                    raise NotImplementedError
+            self.interface_blocks = nn.Sequential(*interface_blocks)
+
+            interface_norms = []
+            for i in range(block_nums):
+                interface_norms.append(norm_layer(embed_dim))
+            self.interface_norms = nn.Sequential(*interface_norms)
+
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+
+        self.instruct = instruct
+        if instruct:
+            num_embeddings = 4
+            self.prompt_embeddings = nn.Embedding(num_embeddings, embed_dim) # should be consistent with new tokens in decoder.instruct_tokens
+
+        self.norm = norm_layer(embed_dim)
+        trunc_normal_(self.pos_embed_search, std=.02)
+        trunc_normal_(self.pos_embed_template, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, template_list, search_list, text_src, seq):
+        num_template = len(template_list)
+        num_search = len(search_list)
+        if self.instruct:
+            instruct_embedding = self.prompt_embeddings(seq).unsqueeze(1)
+
+        z = torch.stack(template_list, dim=1)#(b,n,c,h,w)
+        z = z.view(-1, *z.size()[2:])#(bn,c,h,w)
+        x = torch.stack(search_list, dim=1)#(b,n,c,h,w)
+        x = x.view(-1, *x.size()[2:])#(bn,c,h,w)
+
+        # rgb image
+        x_rgb = x[:, :3, :, :]
+        z_rgb = z[:, :3, :, :]
+        # multi-modal image
+        x_dte = x[:, 3:, :, :]
+        z_dte = z[:, 3:, :, :]
+
+        x_rgb = self.patch_embed(x_rgb)
+        z_rgb = self.patch_embed(z_rgb)
+
+        if self.interface_type in ['low-rank_add']:
+            z_dte = self.patch_embed_interface(z_dte)
+            x_dte = self.patch_embed_interface(x_dte)
+            z_dte, x_dte = self.language_interface(z_dte, x_dte, text_src) # add language information
+            z_rgb_feat = token2feature(self.interface_norms[0](z_rgb))
+            x_rgb_feat = token2feature(self.interface_norms[0](x_rgb))
+            z_dte_feat = token2feature(self.interface_norms[0](z_dte))
+            x_dte_feat = token2feature(self.interface_norms[0](x_dte))
+            z_feat = torch.cat([z_rgb_feat, z_dte_feat], dim=1)
+            x_feat = torch.cat([x_rgb_feat, x_dte_feat], dim=1)
+            z_feat = self.interface_blocks[0](z_feat)
+            x_feat = self.interface_blocks[0](x_feat)
+            z_dte = feature2token(z_feat)
+            x_dte = feature2token(x_feat)
+            x = x_rgb + x_dte
+            z = z_rgb + z_dte
+            z = z + self.pos_embed_template
+            x = x + self.pos_embed_search
+            x_dte = x_dte.reshape(-1, num_search * x_dte.size(1), x_dte.size(-1))
+            z_dte = z_dte.reshape(-1, num_template * z_dte.size(1), z_dte.size(-1))
+            z = z.reshape(-1, num_template * z.size(1), z.size(-1))
+            x = x.reshape(-1, num_search * x.size(1), x.size(-1))
+            len_x = x.size(1)
+            len_z = z.size(1)
+            xz = torch.cat([x, z], dim=1)
+        else:
+            raise ValueError('illegal interface_type')
+
+        if self.instruct:
+            xz = torch.cat([instruct_embedding, xz], dim=1)
+
+        xz = self.pos_drop(xz)
+
+        for i, blk in enumerate(self.blocks):   #batch is the first dimension.
+            if i >= 1:
+                if self.interface_type in ['low-rank_add']:
+                    if self.instruct:
+                        instruct_embedding = xz[:, 0, :].unsqueeze(1)
+                        xz = xz[:, 1:, :]
+                    xz_ori = xz
+                    x = xz[:, :len_x, :]
+                    z = xz[:, len_x:, :]
+                    x = x.reshape(x.size(0)*num_search,-1,x.size(-1))
+                    z = z.reshape(z.size(0)*num_template,-1,z.size(-1))
+                    x_dte = x_dte.reshape(x_dte.size(0)*num_search,-1,x.size(-1))
+                    z_dte = z_dte.reshape(z_dte.size(0)*num_template,-1,z.size(-1))
+                    x_rgb_feat = token2feature(self.interface_norms[i](x))
+                    z_rgb_feat = token2feature(self.interface_norms[i](z))
+                    x_dte_feat = token2feature(self.interface_norms[i](x_dte))
+                    z_dte_feat = token2feature(self.interface_norms[i](z_dte))
+                    z_feat = torch.cat([z_rgb_feat, z_dte_feat], dim=1)
+                    x_feat = torch.cat([x_rgb_feat, x_dte_feat], dim=1)
+                    z_feat = self.interface_blocks[i](z_feat)
+                    x_feat = self.interface_blocks[i](x_feat)
+                    z_dte = feature2token(z_feat)
+                    x_dte = feature2token(x_feat)
+                    x_dte = x_dte.reshape(-1,num_search*x_dte.size(1),x_dte.size(-1))
+                    z_dte = z_dte.reshape(-1,num_template*z_dte.size(1),z_dte.size(-1))
+                    xz_dte = torch.cat([x_dte, z_dte], dim=1)
+                    xz = xz_ori + xz_dte
+                    if self.instruct:
+                        xz = torch.cat([instruct_embedding, xz], dim=1)
+
+            if self.use_checkpoint:
+                xz = checkpoint.checkpoint(blk, xz)
+            else:
+                xz = blk(xz)
+
+        xz = self.norm(xz) # B,N,C
+        return xz
+
+    def forward_features_rgb(self, template_list, search_list):
+        num_template = len(template_list)
+        num_search = len(search_list)
+
+        z = torch.stack(template_list, dim=1)#(b,n,c,h,w)
+        z = z.view(-1, *z.size()[2:])#(bn,c,h,w)
+        x = torch.stack(search_list, dim=1)#(b,n,c,h,w)
+        x = x.view(-1, *x.size()[2:])#(bn,c,h,w)
+
+        x = self.patch_embed(x)
+        z = self.patch_embed(z)
+
+
+        z = z + self.pos_embed_template
+        x = x + self.pos_embed_search
+
+        # for multiple search region and template, go back
+        z = z.reshape(-1,num_template * z.size(1),z.size(-1))
+        x = x.reshape(-1,num_search * x.size(1),x.size(-1))
+
+        len_x = x.size(1)
+        len_z = z.size(1)
+
+        xz = torch.cat([x, z], dim=1)
+
+        xz = self.pos_drop(xz)
+
+        for i, blk in enumerate(self.blocks):   #batch is the first dimension.
+            if self.use_checkpoint:
+                xz = checkpoint.checkpoint(blk, xz)
+            else:
+                print(i)
+                xz = blk(xz)
+
+        xz = self.norm(xz) # B,N,C
+        return xz
+
+    def forward(self, template_list, search_list, text_src, seq):
+        xz = self.forward_features(template_list, search_list, text_src, seq)
+        out=[xz]
+        return out
+
+    def forward_rgb(self, template_list, search_list):
+        xz = self.forward_features_rgb(template_list, search_list)
+        out=[xz]
+        return out
+
+    def language_interface(self, z_dte, x_dte, text_src):
+        text_src = text_src.unsqueeze(1)
+        x_dte = x_dte * text_src
+        text_src_z = text_src.expand(-1,self.num_template,-1).reshape(text_src.size(0)*self.num_template,1,-1)
+        z_dte = z_dte * text_src_z
+        return z_dte, x_dte
+
+@register_model
+def vitmm_base_patch16(pretrained=False, pretrain_type='default',
+                       search_size=384, template_size=192, **kwargs):
+    patch_size = 16
+    model = VisionTransformerMM(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_base_patch16_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model, pretrain_type,
+                        num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+@register_model
+def vitmm_large_patch16(pretrained=False, pretrain_type='default',
+                        search_size=384, template_size=192, **kwargs):
+    patch_size = 16
+    model = VisionTransformerMM(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_large_patch16_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model, pretrain_type, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+@register_model
+def vitmm_huge_patch14(pretrained=False, pretrain_type='default',
+                       search_size=364, template_size=182, **kwargs):
+    patch_size = 14
+    model = VisionTransformerMM(
+        search_size=search_size, template_size=template_size,
+        patch_size=patch_size, num_classes=0,
+        embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    cfg_type = 'vit_huge_patch14_224_' + pretrain_type
+    if pretrain_type == 'scratch':
+        pretrained = False
+        return model
+    model.default_cfg = default_cfgs[cfg_type]
+    if pretrained:
+        load_pretrained(model,
+                        pretrain_type, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model
+
+def load_pretrained(model, pretrain_type='default', cfg=None, num_classes=1000, in_chans=3, filter_fn=None, strict=False):
+    if cfg is None:
+        cfg = getattr(model, 'default_cfg')
+    if cfg is None or 'url' not in cfg or not cfg['url']:
+        print("Pretrained model URL is invalid, using random initialization.")
+        return
+
+    state_dict = model_zoo.load_url(cfg['url'], progress=False, map_location='cpu')
+    if pretrain_type == 'mae':
+        state_dict = state_dict['model']
+
+    if filter_fn is not None:
+        state_dict = filter_fn(state_dict)
+
+    if in_chans == 1:
+        conv1_name = cfg['first_conv']
+        print('Converting first conv (%s) pretrained weights from 3 to 1 channel' % conv1_name)
+        conv1_weight = state_dict[conv1_name + '.weight']
+        # Some weights are in torch.half, ensure it's float for sum on CPU
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I > 3:
+            assert conv1_weight.shape[1] % 3 == 0
+            # For models with space2depth stems
+            conv1_weight = conv1_weight.reshape(O, I // 3, 3, J, K)
+            conv1_weight = conv1_weight.sum(dim=2, keepdim=False)
+        else:
+            conv1_weight = conv1_weight.sum(dim=1, keepdim=True)
+        conv1_weight = conv1_weight.to(conv1_type)
+        state_dict[conv1_name + '.weight'] = conv1_weight
+    elif in_chans != 3:
+        conv1_name = cfg['first_conv']
+        conv1_weight = state_dict[conv1_name + '.weight']
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I != 3:
+            print('Deleting first conv (%s) from pretrained weights.' % conv1_name)
+            del state_dict[conv1_name + '.weight']
+            strict = False
+        else:
+            # NOTE this strategy should be better than random init, but there could be other combinations of
+            # the original RGB input layer weights that'd work better for specific cases.
+            print('Repeating first conv (%s) weights in channel dim.' % conv1_name)
+            repeat = int(math.ceil(in_chans / 3))
+            conv1_weight = conv1_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
+            conv1_weight *= (3 / float(in_chans))
+            conv1_weight = conv1_weight.to(conv1_type)
+            state_dict[conv1_name + '.weight'] = conv1_weight
+
+    classifier_name = cfg['classifier']
+    if pretrain_type == "mae":
+        pass
+    elif num_classes == 1000 and cfg['num_classes'] == 1001:
+        # special case for imagenet trained models with extra background class in pretrained weights
+        classifier_weight = state_dict[classifier_name + '.weight']
+        state_dict[classifier_name + '.weight'] = classifier_weight[1:]
+        classifier_bias = state_dict[classifier_name + '.bias']
+        state_dict[classifier_name + '.bias'] = classifier_bias[1:]
+    elif num_classes != cfg['num_classes']:
+        # completely discard fully connected for all other differences between pretrained and created model
+        del state_dict[classifier_name + '.weight']
+        del state_dict[classifier_name + '.bias']
+
+    # adjust position encoding
+    pe = state_dict['pos_embed'][:,1:,:]
+    b_pe, hw_pe, c_pe = pe.shape
+    side_pe = int(math.sqrt(hw_pe))
+    side_num_patches_search = int(math.sqrt(model.num_patches_search))
+    side_num_patches_template = int(math.sqrt(model.num_patches_template))
+    pe_2D = pe.reshape([b_pe, side_pe, side_pe, c_pe]).permute([0,3,1,2])  #b,c,h,w
+    if side_pe != side_num_patches_search:
+        pe_s_2D = nn.functional.interpolate(pe_2D, [side_num_patches_search, side_num_patches_search], align_corners=True, mode='bicubic')
+        pe_s = torch.flatten(pe_s_2D.permute([0,2,3,1]),1,2)
+    else:
+        pe_s = pe
+    if side_pe != side_num_patches_template:
+        pe_t_2D = nn.functional.interpolate(pe_2D, [side_num_patches_template, side_num_patches_template], align_corners=True, mode='bicubic')
+        pe_t = torch.flatten(pe_t_2D.permute([0, 2, 3, 1]), 1, 2)
+    else:
+        pe_t = pe
+    state_dict['pos_embed_template'] = pe_t
+    state_dict['pos_embed_search'] = pe_s
+    del state_dict['cls_token']
+    del state_dict['pos_embed']
+
+    model.load_state_dict(state_dict, strict=strict)
+

lib/models/layers/adapter.py

@@ -0,0 +1,104 @@
+import torch
+from torch import nn
+import timm
+import math
+
+
+'''
+def forward_block(self, x):
+    x = x + self.drop_path(self.attn(self.norm1(x))) + self.drop_path(self.adapter_attn(self.norm1(x))) * self.s
+    x = x + self.drop_path(self.mlp(self.norm2(x))) + self.drop_path(self.adapter_mlp(self.norm2(x))) * self.s
+    return x
+
+
+def forward_block_attn(self, x):
+    x = x + self.drop_path(self.attn(self.norm1(x))) + self.drop_path(self.adapter_attn(self.norm1(x))) * self.s
+    x = x + self.drop_path(self.mlp(self.norm2(x)))
+    return x
+'''
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+
+class Bi_direct_adapter(nn.Module):
+    def __init__(self, dim=8,input_dim=768, xavier_init=False):
+        super().__init__()
+
+        self.adapter_down = nn.Linear(input_dim, dim)
+        self.adapter_up = nn.Linear(dim, input_dim)
+        self.adapter_mid = nn.Linear(dim, dim)
+
+        #nn.init.xavier_uniform_(self.adapter_down.weight)
+        nn.init.zeros_(self.adapter_mid.bias)
+        nn.init.zeros_(self.adapter_mid.weight)
+        nn.init.zeros_(self.adapter_down.weight)
+        nn.init.zeros_(self.adapter_down.bias)
+        nn.init.zeros_(self.adapter_up.weight)
+        nn.init.zeros_(self.adapter_up.bias)
+
+        #self.act = QuickGELU()
+        self.dropout = nn.Dropout(0.1)
+        self.dim = dim
+
+    def forward(self, x):
+        # B, N, C = x.shape
+        x_down = self.adapter_down(x)   
+        #x_down = self.act(x_down)
+        x_down = self.adapter_mid(x_down)
+        #x_down = self.act(x_down)
+        x_down = self.dropout(x_down)
+        x_up = self.adapter_up(x_down)  
+        #print("return adap x", x_up.size())
+        return x_up
+
+"""
+
+
+class Convpass(nn.Module):
+    def __init__(self, dim=8, xavier_init=False):
+        super().__init__()
+
+        self.adapter_conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        if xavier_init:
+            nn.init.xavier_uniform_(self.adapter_conv.weight)
+        else:
+            nn.init.zeros_(self.adapter_conv.weight)
+            self.adapter_conv.weight.data[:, :, 1, 1] += torch.eye(8, dtype=torch.float)
+        nn.init.zeros_(self.adapter_conv.bias)
+
+        self.adapter_down = nn.Linear(768, dim)  # equivalent to 1 * 1 Conv
+        self.adapter_up = nn.Linear(dim, 768)  # equivalent to 1 * 1 Conv
+        nn.init.xavier_uniform_(self.adapter_down.weight)
+        nn.init.zeros_(self.adapter_down.bias)
+        nn.init.zeros_(self.adapter_up.weight)
+        nn.init.zeros_(self.adapter_up.bias)
+
+        self.act = QuickGELU()
+        self.dropout = nn.Dropout(0.1)
+        self.dim = dim
+
+    def forward(self, x):
+        B, N, C = x.shape
+        #print(x.shape)
+        x_down = self.adapter_down(x)  # equivalent to 1 * 1 Conv
+        x_down = self.act(x_down)
+
+        #print(x_down.shape)
+
+        x_patch = x_down[:, 64:].reshape(B, 16, 16, self.dim).permute(0, 3, 1, 2)   ############
+        x_patch = self.adapter_conv(x_patch)
+        x_patch = x_patch.permute(0, 2, 3, 1).reshape(B, 16 * 16, self.dim)
+
+
+        #x_down = torch.cat([x_cls, x_patch], dim=1)
+
+        x_down = self.act(x_down)
+        x_down = self.dropout(x_down)
+        x_up = self.adapter_up(x_down)  # equivalent to 1 * 1 Conv
+
+        return x_up
+"""

lib/models/layers/attn.py

@@ -0,0 +1,128 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_
+
+from lib.models.layers.rpe import generate_2d_concatenated_self_attention_relative_positional_encoding_index
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.,
+                 rpe=False, z_size=7, x_size=14):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.rpe =rpe
+        if self.rpe:
+            relative_position_index = \
+                generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
+                                                                                           [x_size, x_size])
+            self.register_buffer("relative_position_index", relative_position_index)
+            # define a parameter table of relative position bias
+            self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
+                                                                          relative_position_index.max() + 1)))
+            trunc_normal_(self.relative_position_bias_table, std=0.02)
+
+    def forward(self, x, mask=None, return_attention=False):
+        # x: B, N, C
+        # mask: [B, N, ] torch.bool
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        if self.rpe:
+            relative_position_bias = self.relative_position_bias_table[:, self.relative_position_index].unsqueeze(0)
+            attn += relative_position_bias
+
+        if mask is not None:
+            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2), float('-inf'),)
+
+        split_attn = False
+        len_t = 49
+        if split_attn:
+            attn_t = attn[..., :len_t].softmax(dim=-1)
+            attn_s = attn[..., len_t:].softmax(dim=-1)
+            attn = torch.cat([attn_t, attn_s], dim=-1)
+        else:
+            attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        if return_attention:
+            return x, attn
+        else:
+            return x
+
+
+class Attention_talking_head(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    # with slight modifications to add Talking Heads Attention (https://arxiv.org/pdf/2003.02436v1.pdf)
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
+                 rpe=True, z_size=7, x_size=14):
+        super().__init__()
+
+        self.num_heads = num_heads
+
+        head_dim = dim // num_heads
+
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+
+        self.proj = nn.Linear(dim, dim)
+
+        self.proj_l = nn.Linear(num_heads, num_heads)
+        self.proj_w = nn.Linear(num_heads, num_heads)
+
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.rpe = rpe
+        if self.rpe:
+            relative_position_index = \
+                generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
+                                                                                           [x_size, x_size])
+            self.register_buffer("relative_position_index", relative_position_index)
+            # define a parameter table of relative position bias
+            self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
+                                                                          relative_position_index.max() + 1)))
+            trunc_normal_(self.relative_position_bias_table, std=0.02)
+
+    def forward(self, x, mask=None):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1))
+
+        if self.rpe:
+            relative_position_bias = self.relative_position_bias_table[:, self.relative_position_index].unsqueeze(0)
+            attn += relative_position_bias
+
+        if mask is not None:
+            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2),
+                                    float('-inf'),)
+
+        attn = self.proj_l(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+        attn = attn.softmax(dim=-1)
+
+        attn = self.proj_w(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

lib/models/layers/attn_adapt_blocks.py

@@ -0,0 +1,155 @@
+import math
+import torch
+import torch.nn as nn
+from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+
+from lib.models.layers.attn import Attention
+from lib.models.layers.adapter import Bi_direct_adapter
+
+
+
+def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor, lens_t: int, keep_ratio: float, global_index: torch.Tensor, box_mask_z: torch.Tensor):
+    """
+    Eliminate potential background candidates for computation reduction and noise cancellation.
+    Args:
+        attn (torch.Tensor): [B, num_heads, L_t + L_s, L_t + L_s], attention weights
+        tokens (torch.Tensor):  [B, L_t + L_s, C], template and search region tokens
+        lens_t (int): length of template
+        keep_ratio (float): keep ratio of search region tokens (candidates)
+        global_index (torch.Tensor): global index of search region tokens
+        box_mask_z (torch.Tensor): template mask used to accumulate attention weights
+
+    Returns:
+        tokens_new (torch.Tensor): tokens after candidate elimination
+        keep_index (torch.Tensor): indices of kept search region tokens
+        removed_index (torch.Tensor): indices of removed search region tokens
+    """
+    lens_s = attn.shape[-1] - lens_t    
+    bs, hn, _, _ = attn.shape
+
+    lens_keep = math.ceil(keep_ratio * lens_s)
+    if lens_keep == lens_s:
+        return tokens, global_index, None
+
+    attn_t = attn[:, :, :lens_t, lens_t:]
+
+    
+
+
+    if box_mask_z is not None:
+        #print("\n1\n1\n1")
+        box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(-1, attn_t.shape[1], -1, attn_t.shape[-1])
+        # attn_t = attn_t[:, :, box_mask_z, :]
+        attn_t = attn_t[box_mask_z]
+        attn_t = attn_t.view(bs, hn, -1, lens_s)
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+        # attn_t = [attn_t[i, :, box_mask_z[i, :], :] for i in range(attn_t.size(0))]
+        # attn_t = [attn_t[i].mean(dim=1).mean(dim=0) for i in range(len(attn_t))]
+        # attn_t = torch.stack(attn_t, dim=0)
+    else:
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+    # use sort instead of topk, due to the speed issue
+    # https://github.com/pytorch/pytorch/issues/22812
+    sorted_attn, indices = torch.sort(attn_t, dim=1, descending=True)
+
+
+
+    topk_attn, topk_idx = sorted_attn[:, :lens_keep], indices[:, :lens_keep]
+    non_topk_attn, non_topk_idx = sorted_attn[:, lens_keep:], indices[:, lens_keep:]
+    
+    keep_index = global_index.gather(dim=1, index=topk_idx)
+    
+    removed_index = global_index.gather(dim=1, index=non_topk_idx)
+    
+
+    # separate template and search tokens
+    tokens_t = tokens[:, :lens_t]
+    tokens_s = tokens[:, lens_t:]
+
+    # obtain the attentive and inattentive tokens
+    B, L, C = tokens_s.shape
+    # topk_idx_ = topk_idx.unsqueeze(-1).expand(B, lens_keep, C)
+
+    attentive_tokens = tokens_s.gather(dim=1, index=topk_idx.unsqueeze(-1).expand(B, -1, C))
+    # inattentive_tokens = tokens_s.gather(dim=1, index=non_topk_idx.unsqueeze(-1).expand(B, -1, C))
+
+    # compute the weighted combination of inattentive tokens
+    # fused_token = non_topk_attn @ inattentive_tokens
+    
+    # concatenate these tokens
+    # tokens_new = torch.cat([tokens_t, attentive_tokens, fused_token], dim=0)
+    tokens_new = torch.cat([tokens_t, attentive_tokens], dim=1)
+
+    #print("finish ce func")
+
+    return tokens_new, keep_index, removed_index                       # x, global_index_search, removed_index_search
+
+
+class CEABlock(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, keep_ratio_search=1.0,):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)     #from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+
+        self.keep_ratio_search = keep_ratio_search
+
+
+        self.adap_t = Bi_direct_adapter()        
+        self.adap2_t = Bi_direct_adapter()
+
+    def forward(self, x, xi, global_index_template, global_index_templatei, global_index_search, global_index_searchi, mask=None, ce_template_mask=None, keep_ratio_search=None):
+        
+        xori = x
+        
+        x_attn, attn = self.attn(self.norm1(x), mask, True)   
+        x = x + self.drop_path(x_attn) + self.drop_path(self.adap_t(self.norm1(xi)))  #########-------------------------adapter
+
+        xi_attn, i_attn = self.attn(self.norm1(xi), mask,True)
+        xi = xi + self.drop_path(xi_attn) + self.drop_path(self.adap_t(self.norm1(xori)))  #########-------------------------adapter
+                     
+        lens_t = global_index_template.shape[1]
+
+        removed_index_search = None
+        removed_index_searchi = None
+        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
+            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
+            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)
+            xi, global_index_searchi, removed_index_searchi = candidate_elimination(i_attn, xi, lens_t, keep_ratio_search, global_index_searchi, ce_template_mask)
+
+        xori = x
+
+        x = x + self.drop_path(self.mlp(self.norm2(x))) + self.drop_path(self.adap2_t(self.norm2(xi)))   ###-------adapter
+
+        xi = xi + self.drop_path(self.mlp(self.norm2(xi))) + self.drop_path(self.adap2_t(self.norm2(xori)))   ###-------adapter
+        
+        return x, global_index_template, global_index_search, removed_index_search, attn, xi, global_index_templatei, global_index_searchi, removed_index_searchi, i_attn
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        #print("class Block ")
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, mask=None):
+        #print("class Block forward")
+        x = x + self.drop_path(self.attn(self.norm1(x), mask))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x

lib/models/layers/attn_blocks.py

@@ -0,0 +1,122 @@
+import math
+import torch
+import torch.nn as nn
+from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+
+from lib.models.layers.attn import Attention
+
+
+def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor, lens_t: int, keep_ratio: float, global_index: torch.Tensor, box_mask_z: torch.Tensor):
+    """
+    Eliminate potential background candidates for computation reduction and noise cancellation.
+    Args:
+        attn (torch.Tensor): [B, num_heads, L_t + L_s, L_t + L_s], attention weights
+        tokens (torch.Tensor):  [B, L_t + L_s, C], template and search region tokens
+        lens_t (int): length of template
+        keep_ratio (float): keep ratio of search region tokens (candidates)
+        global_index (torch.Tensor): global index of search region tokens
+        box_mask_z (torch.Tensor): template mask used to accumulate attention weights
+
+    Returns:
+        tokens_new (torch.Tensor): tokens after candidate elimination
+        keep_index (torch.Tensor): indices of kept search region tokens
+        removed_index (torch.Tensor): indices of removed search region tokens
+    """
+    lens_s = attn.shape[-1] - lens_t
+    bs, hn, _, _ = attn.shape
+
+    lens_keep = math.ceil(keep_ratio * lens_s)
+    if lens_keep == lens_s:
+        return tokens, global_index, None
+
+    attn_t = attn[:, :, :lens_t, lens_t:]
+
+    if box_mask_z is not None:
+        box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(-1, attn_t.shape[1], -1, attn_t.shape[-1])
+        # attn_t = attn_t[:, :, box_mask_z, :]
+        attn_t = attn_t[box_mask_z]
+        attn_t = attn_t.view(bs, hn, -1, lens_s)
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+        # attn_t = [attn_t[i, :, box_mask_z[i, :], :] for i in range(attn_t.size(0))]
+        # attn_t = [attn_t[i].mean(dim=1).mean(dim=0) for i in range(len(attn_t))]
+        # attn_t = torch.stack(attn_t, dim=0)
+    else:
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+    # use sort instead of topk, due to the speed issue
+    # https://github.com/pytorch/pytorch/issues/22812
+    sorted_attn, indices = torch.sort(attn_t, dim=1, descending=True)
+
+    topk_attn, topk_idx = sorted_attn[:, :lens_keep], indices[:, :lens_keep]
+    non_topk_attn, non_topk_idx = sorted_attn[:, lens_keep:], indices[:, lens_keep:]
+
+    keep_index = global_index.gather(dim=1, index=topk_idx)
+    removed_index = global_index.gather(dim=1, index=non_topk_idx)
+
+    # separate template and search tokens
+    tokens_t = tokens[:, :lens_t]
+    tokens_s = tokens[:, lens_t:]
+
+    # obtain the attentive and inattentive tokens
+    B, L, C = tokens_s.shape
+    # topk_idx_ = topk_idx.unsqueeze(-1).expand(B, lens_keep, C)
+    attentive_tokens = tokens_s.gather(dim=1, index=topk_idx.unsqueeze(-1).expand(B, -1, C))
+    # inattentive_tokens = tokens_s.gather(dim=1, index=non_topk_idx.unsqueeze(-1).expand(B, -1, C))
+
+    # compute the weighted combination of inattentive tokens
+    # fused_token = non_topk_attn @ inattentive_tokens
+
+    # concatenate these tokens
+    # tokens_new = torch.cat([tokens_t, attentive_tokens, fused_token], dim=0)
+    tokens_new = torch.cat([tokens_t, attentive_tokens], dim=1)
+
+    return tokens_new, keep_index, removed_index
+
+
+class CEBlock(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, keep_ratio_search=1.0,):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        self.keep_ratio_search = keep_ratio_search
+
+    def forward(self, x, global_index_template, global_index_search, mask=None, ce_template_mask=None, keep_ratio_search=None):
+        x_attn, attn = self.attn(self.norm1(x), mask, True)
+        x = x + self.drop_path(x_attn)
+        lens_t = global_index_template.shape[1]
+
+        removed_index_search = None
+        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
+            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
+            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)
+
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x, global_index_template, global_index_search, removed_index_search, attn
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, mask=None):
+        x = x + self.drop_path(self.attn(self.norm1(x), mask))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x

lib/models/layers/dualstream_attn_blocks.py

@@ -0,0 +1,152 @@
+import math
+import torch
+import torch.nn as nn
+from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+
+from lib.models.layers.attn import Attention
+from lib.models.layers.adapter import Bi_direct_adapter
+
+
+
+
+
+def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor, lens_t: int, keep_ratio: float, global_index: torch.Tensor, box_mask_z: torch.Tensor):
+    """
+    Eliminate potential background candidates for computation reduction and noise cancellation.
+    Args:
+        attn (torch.Tensor): [B, num_heads, L_t + L_s, L_t + L_s], attention weights
+        tokens (torch.Tensor):  [B, L_t + L_s, C], template and search region tokens
+        lens_t (int): length of template
+        keep_ratio (float): keep ratio of search region tokens (candidates)
+        global_index (torch.Tensor): global index of search region tokens
+        box_mask_z (torch.Tensor): template mask used to accumulate attention weights
+
+    Returns:
+        tokens_new (torch.Tensor): tokens after candidate elimination
+        keep_index (torch.Tensor): indices of kept search region tokens
+        removed_index (torch.Tensor): indices of removed search region tokens
+    """
+    lens_s = attn.shape[-1] - lens_t    
+    bs, hn, _, _ = attn.shape
+
+    lens_keep = math.ceil(keep_ratio * lens_s)
+    if lens_keep == lens_s:
+        return tokens, global_index, None
+
+    attn_t = attn[:, :, :lens_t, lens_t:]
+
+    #print("box_mask_z")
+
+
+    if box_mask_z is not None:
+        #print("\n1\n1\n1")
+        box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(-1, attn_t.shape[1], -1, attn_t.shape[-1])
+        # attn_t = attn_t[:, :, box_mask_z, :]
+        attn_t = attn_t[box_mask_z]
+        attn_t = attn_t.view(bs, hn, -1, lens_s)
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+        # attn_t = [attn_t[i, :, box_mask_z[i, :], :] for i in range(attn_t.size(0))]
+        # attn_t = [attn_t[i].mean(dim=1).mean(dim=0) for i in range(len(attn_t))]
+        # attn_t = torch.stack(attn_t, dim=0)
+    else:
+        attn_t = attn_t.mean(dim=2).mean(dim=1)  # B, H, L-T, L_s --> B, L_s
+
+    # use sort instead of topk, due to the speed issue
+    # https://github.com/pytorch/pytorch/issues/22812
+    sorted_attn, indices = torch.sort(attn_t, dim=1, descending=True)
+
+
+
+    topk_attn, topk_idx = sorted_attn[:, :lens_keep], indices[:, :lens_keep]
+    non_topk_attn, non_topk_idx = sorted_attn[:, lens_keep:], indices[:, lens_keep:]
+    
+    keep_index = global_index.gather(dim=1, index=topk_idx)
+    
+    removed_index = global_index.gather(dim=1, index=non_topk_idx)
+    
+
+    # separate template and search tokens
+    tokens_t = tokens[:, :lens_t]
+    tokens_s = tokens[:, lens_t:]
+
+    # obtain the attentive and inattentive tokens
+    B, L, C = tokens_s.shape
+    # topk_idx_ = topk_idx.unsqueeze(-1).expand(B, lens_keep, C)
+
+    attentive_tokens = tokens_s.gather(dim=1, index=topk_idx.unsqueeze(-1).expand(B, -1, C))
+    # inattentive_tokens = tokens_s.gather(dim=1, index=non_topk_idx.unsqueeze(-1).expand(B, -1, C))
+
+    # compute the weighted combination of inattentive tokens
+    # fused_token = non_topk_attn @ inattentive_tokens
+    
+    # concatenate these tokens
+    # tokens_new = torch.cat([tokens_t, attentive_tokens, fused_token], dim=0)
+    tokens_new = torch.cat([tokens_t, attentive_tokens], dim=1)
+
+    #print("finish ce func")
+
+    return tokens_new, keep_index, removed_index                       # x, global_index_search, removed_index_search
+
+
+class DSBlock(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, keep_ratio_search=1.0,):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)     #from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+
+        self.keep_ratio_search = keep_ratio_search
+
+
+    def forward(self, x, xi, global_index_template, global_index_templatei, global_index_search, global_index_searchi, mask=None, ce_template_mask=None, keep_ratio_search=None):
+        
+        
+        x_attn, attn = self.attn(self.norm1(x), mask, True)   
+        x = x + self.drop_path(x_attn) 
+        xi_attn, i_attn = self.attn(self.norm1(xi), mask,True)
+        xi = xi + self.drop_path(xi_attn) 
+                     
+        lens_t = global_index_template.shape[1]
+
+        removed_index_search = None
+        removed_index_searchi = None
+        if self.keep_ratio_search < 1 and (keep_ratio_search is None or keep_ratio_search < 1):
+            keep_ratio_search = self.keep_ratio_search if keep_ratio_search is None else keep_ratio_search
+            x, global_index_search, removed_index_search = candidate_elimination(attn, x, lens_t, keep_ratio_search, global_index_search, ce_template_mask)
+            xi, global_index_searchi, removed_index_searchi = candidate_elimination(i_attn, xi, lens_t, keep_ratio_search, global_index_searchi, ce_template_mask)
+
+        
+
+        x = x + self.drop_path(self.mlp(self.norm2(x))) 
+
+        xi = xi + self.drop_path(self.mlp(self.norm2(xi))) 
+        
+        return x, global_index_template, global_index_search, removed_index_search, attn, xi, global_index_templatei, global_index_searchi, removed_index_searchi, i_attn
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        #print("class Block ")
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, mask=None):
+        #print("class Block forward")
+        x = x + self.drop_path(self.attn(self.norm1(x), mask))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x

lib/models/layers/frozen_bn.py

@@ -0,0 +1,40 @@
+import torch
+
+
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed.
+
+    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
+    without which any other models than torchvision.models.resnet[18,34,50,101]
+    produce nans.
+    """
+
+    def __init__(self, n):
+        super(FrozenBatchNorm2d, self).__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        num_batches_tracked_key = prefix + 'num_batches_tracked'
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+
+        super(FrozenBatchNorm2d, self)._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs)
+
+    def forward(self, x):
+        # move reshapes to the beginning
+        # to make it fuser-friendly
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        eps = 1e-5
+        scale = w * (rv + eps).rsqrt()  # rsqrt(x): 1/sqrt(x), r: reciprocal
+        bias = b - rm * scale
+        return x * scale + bias

lib/models/layers/head.py

@@ -0,0 +1,248 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+from lib.models.layers.frozen_bn import FrozenBatchNorm2d
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1,
+         freeze_bn=False):
+    if freeze_bn:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            FrozenBatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+    else:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+
+
+class Corner_Predictor(nn.Module):
+    """ Corner Predictor module"""
+
+    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
+        super(Corner_Predictor, self).__init__()
+        self.feat_sz = feat_sz
+        self.stride = stride
+        self.img_sz = self.feat_sz * self.stride
+        '''top-left corner'''
+        self.conv1_tl = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_tl = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_tl = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_tl = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_tl = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        '''bottom-right corner'''
+        self.conv1_br = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_br = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_br = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_br = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_br = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        '''about coordinates and indexs'''
+        with torch.no_grad():
+            self.indice = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride
+            # generate mesh-grid
+            self.coord_x = self.indice.repeat((self.feat_sz, 1)) \
+                .view((self.feat_sz * self.feat_sz,)).float().cuda()
+            self.coord_y = self.indice.repeat((1, self.feat_sz)) \
+                .view((self.feat_sz * self.feat_sz,)).float().cuda()
+
+    def forward(self, x, return_dist=False, softmax=True):
+        """ Forward pass with input x. """
+        score_map_tl, score_map_br = self.get_score_map(x)
+        if return_dist:
+            coorx_tl, coory_tl, prob_vec_tl = self.soft_argmax(score_map_tl, return_dist=True, softmax=softmax)
+            coorx_br, coory_br, prob_vec_br = self.soft_argmax(score_map_br, return_dist=True, softmax=softmax)
+            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz, prob_vec_tl, prob_vec_br
+        else:
+            coorx_tl, coory_tl = self.soft_argmax(score_map_tl)
+            coorx_br, coory_br = self.soft_argmax(score_map_br)
+            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz
+
+    def get_score_map(self, x):
+        # top-left branch
+        x_tl1 = self.conv1_tl(x)
+        x_tl2 = self.conv2_tl(x_tl1)
+        x_tl3 = self.conv3_tl(x_tl2)
+        x_tl4 = self.conv4_tl(x_tl3)
+        score_map_tl = self.conv5_tl(x_tl4)
+
+        # bottom-right branch
+        x_br1 = self.conv1_br(x)
+        x_br2 = self.conv2_br(x_br1)
+        x_br3 = self.conv3_br(x_br2)
+        x_br4 = self.conv4_br(x_br3)
+        score_map_br = self.conv5_br(x_br4)
+        return score_map_tl, score_map_br
+
+    def soft_argmax(self, score_map, return_dist=False, softmax=True):
+        """ get soft-argmax coordinate for a given heatmap """
+        score_vec = score_map.view((-1, self.feat_sz * self.feat_sz))  # (batch, feat_sz * feat_sz)
+        prob_vec = nn.functional.softmax(score_vec, dim=1)
+        exp_x = torch.sum((self.coord_x * prob_vec), dim=1)
+        exp_y = torch.sum((self.coord_y * prob_vec), dim=1)
+        if return_dist:
+            if softmax:
+                return exp_x, exp_y, prob_vec
+            else:
+                return exp_x, exp_y, score_vec
+        else:
+            return exp_x, exp_y
+
+
+class CenterPredictor(nn.Module, ):
+    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
+        super(CenterPredictor, self).__init__()
+        self.feat_sz = feat_sz
+        self.stride = stride
+        self.img_sz = self.feat_sz * self.stride
+
+        # corner predict
+        self.conv1_ctr = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_ctr = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_ctr = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_ctr = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_ctr = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        # size regress
+        self.conv1_offset = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_offset = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_offset = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_offset = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_offset = nn.Conv2d(channel // 8, 2, kernel_size=1)
+
+        # size regress
+        self.conv1_size = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_size = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_size = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_size = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_size = nn.Conv2d(channel // 8, 2, kernel_size=1)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x, gt_score_map=None):
+        """ Forward pass with input x. """
+        score_map_ctr, size_map, offset_map = self.get_score_map(x)
+
+        # assert gt_score_map is None
+        if gt_score_map is None:
+            bbox = self.cal_bbox(score_map_ctr, size_map, offset_map)
+        else:
+            bbox = self.cal_bbox(gt_score_map.unsqueeze(1), size_map, offset_map)
+
+        return score_map_ctr, bbox, size_map, offset_map
+
+    def cal_bbox(self, score_map_ctr, size_map, offset_map, return_score=False):
+        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
+        idx_y = idx // self.feat_sz
+        idx_x = idx % self.feat_sz
+
+        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
+        size = size_map.flatten(2).gather(dim=2, index=idx)
+        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+
+        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
+        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
+        # cx, cy, w, h
+        bbox = torch.cat([(idx_x.to(torch.float) + offset[:, :1]) / self.feat_sz,
+                          (idx_y.to(torch.float) + offset[:, 1:]) / self.feat_sz,
+                          size.squeeze(-1)], dim=1)
+
+        if return_score:
+            return bbox, max_score
+        return bbox
+
+    def get_pred(self, score_map_ctr, size_map, offset_map):
+        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
+        idx_y = idx // self.feat_sz
+        idx_x = idx % self.feat_sz
+
+        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
+        size = size_map.flatten(2).gather(dim=2, index=idx)
+        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+
+        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
+        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
+        return size * self.feat_sz, offset
+
+    def get_score_map(self, x):
+
+        def _sigmoid(x):
+            y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4)
+            return y
+
+        # ctr branch
+        x_ctr1 = self.conv1_ctr(x)
+        x_ctr2 = self.conv2_ctr(x_ctr1)
+        x_ctr3 = self.conv3_ctr(x_ctr2)
+        x_ctr4 = self.conv4_ctr(x_ctr3)
+        score_map_ctr = self.conv5_ctr(x_ctr4)
+
+        # offset branch
+        x_offset1 = self.conv1_offset(x)
+        x_offset2 = self.conv2_offset(x_offset1)
+        x_offset3 = self.conv3_offset(x_offset2)
+        x_offset4 = self.conv4_offset(x_offset3)
+        score_map_offset = self.conv5_offset(x_offset4)
+
+        # size branch
+        x_size1 = self.conv1_size(x)
+        x_size2 = self.conv2_size(x_size1)
+        x_size3 = self.conv3_size(x_size2)
+        x_size4 = self.conv4_size(x_size3)
+        score_map_size = self.conv5_size(x_size4)
+        return _sigmoid(score_map_ctr), _sigmoid(score_map_size), score_map_offset
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, BN=False):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        if BN:
+            self.layers = nn.ModuleList(nn.Sequential(nn.Linear(n, k), nn.BatchNorm1d(k))
+                                        for n, k in zip([input_dim] + h, h + [output_dim]))
+        else:
+            self.layers = nn.ModuleList(nn.Linear(n, k)
+                                        for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def build_box_head(cfg, hidden_dim):
+    stride = cfg.MODEL.BACKBONE.STRIDE
+
+    if cfg.MODEL.HEAD.TYPE == "MLP":
+        mlp_head = MLP(hidden_dim, hidden_dim, 4, 3)  # dim_in, dim_hidden, dim_out, 3 layers
+        return mlp_head
+    elif "CORNER" in cfg.MODEL.HEAD.TYPE:
+        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
+        channel = getattr(cfg.MODEL, "NUM_CHANNELS", 256)
+        print("head channel: %d" % channel)
+        if cfg.MODEL.HEAD.TYPE == "CORNER":
+            corner_head = Corner_Predictor(inplanes=cfg.MODEL.HIDDEN_DIM, channel=channel,
+                                           feat_sz=feat_sz, stride=stride)
+        else:
+            raise ValueError()
+        return corner_head
+    elif cfg.MODEL.HEAD.TYPE == "CENTER":
+        in_channel = hidden_dim
+        out_channel = cfg.MODEL.HEAD.NUM_CHANNELS
+        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
+        center_head = CenterPredictor(inplanes=in_channel, channel=out_channel,
+                                      feat_sz=feat_sz, stride=stride)
+        return center_head
+    else:
+        raise ValueError("HEAD TYPE %s is not supported." % cfg.MODEL.HEAD_TYPE)

lib/models/layers/max_head.py

@@ -0,0 +1,280 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+from lib.models.layers.frozen_bn import FrozenBatchNorm2d
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1,
+         freeze_bn=False):
+    if freeze_bn:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            FrozenBatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+    else:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=True),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU(inplace=True))
+
+
+class Corner_Predictor(nn.Module):
+    """ Corner Predictor module"""
+
+    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
+        super(Corner_Predictor, self).__init__()
+        self.feat_sz = feat_sz
+        self.stride = stride
+        self.img_sz = self.feat_sz * self.stride
+        '''top-left corner'''
+        self.conv1_tl = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_tl = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_tl = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_tl = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_tl = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        '''bottom-right corner'''
+        self.conv1_br = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_br = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_br = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_br = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_br = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        '''about coordinates and indexs'''
+        with torch.no_grad():
+            self.indice = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride
+            # generate mesh-grid
+            self.coord_x = self.indice.repeat((self.feat_sz, 1)) \
+                .view((self.feat_sz * self.feat_sz,)).float().cuda()
+            self.coord_y = self.indice.repeat((1, self.feat_sz)) \
+                .view((self.feat_sz * self.feat_sz,)).float().cuda()
+
+    def forward(self, x, return_dist=False, softmax=True):
+        """ Forward pass with input x. """
+        score_map_tl, score_map_br = self.get_score_map(x)
+        if return_dist:
+            coorx_tl, coory_tl, prob_vec_tl = self.soft_argmax(score_map_tl, return_dist=True, softmax=softmax)
+            coorx_br, coory_br, prob_vec_br = self.soft_argmax(score_map_br, return_dist=True, softmax=softmax)
+            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz, prob_vec_tl, prob_vec_br
+        else:
+            coorx_tl, coory_tl = self.soft_argmax(score_map_tl)
+            coorx_br, coory_br = self.soft_argmax(score_map_br)
+            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz
+
+    def get_score_map(self, x):
+        # top-left branch
+        x_tl1 = self.conv1_tl(x)
+        x_tl2 = self.conv2_tl(x_tl1)
+        x_tl3 = self.conv3_tl(x_tl2)
+        x_tl4 = self.conv4_tl(x_tl3)
+        score_map_tl = self.conv5_tl(x_tl4)
+
+        # bottom-right branch
+        x_br1 = self.conv1_br(x)
+        x_br2 = self.conv2_br(x_br1)
+        x_br3 = self.conv3_br(x_br2)
+        x_br4 = self.conv4_br(x_br3)
+        score_map_br = self.conv5_br(x_br4)
+        return score_map_tl, score_map_br
+
+    def soft_argmax(self, score_map, return_dist=False, softmax=True):
+        """ get soft-argmax coordinate for a given heatmap """
+        score_vec = score_map.view((-1, self.feat_sz * self.feat_sz))  # (batch, feat_sz * feat_sz)
+        prob_vec = nn.functional.softmax(score_vec, dim=1)
+        exp_x = torch.sum((self.coord_x * prob_vec), dim=1)
+        exp_y = torch.sum((self.coord_y * prob_vec), dim=1)
+        if return_dist:
+            if softmax:
+                return exp_x, exp_y, prob_vec
+            else:
+                return exp_x, exp_y, score_vec
+        else:
+            return exp_x, exp_y
+
+
+class CenterPredictor(nn.Module, ):
+    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
+        super(CenterPredictor, self).__init__()
+        self.feat_sz = feat_sz
+        self.stride = stride
+        self.img_sz = self.feat_sz * self.stride
+
+        # corner predict
+        self.conv1_ctr = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_ctr = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_ctr = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_ctr = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_ctr = nn.Conv2d(channel // 8, 1, kernel_size=1)
+
+        # size regress
+        self.conv1_offset = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_offset = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_offset = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_offset = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_offset = nn.Conv2d(channel // 8, 2, kernel_size=1)
+
+        # size regress
+        self.conv1_size = conv(inplanes, channel, freeze_bn=freeze_bn)
+        self.conv2_size = conv(channel, channel // 2, freeze_bn=freeze_bn)
+        self.conv3_size = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
+        self.conv4_size = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
+        self.conv5_size = nn.Conv2d(channel // 8, 2, kernel_size=1)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x, gt_score_map=None):
+        """ Forward pass with input x. """
+        #print("input x shape ", x.shape)
+        score_map_ctr, size_map, offset_map = self.get_score_map(x)
+        #print("score_map_ctr shape  ",score_map_ctr.shape)
+        #print("size_map shape  ", size_map.shape)
+        #print("offset_map shape  ", offset_map.shape)   
+        score_map_ctr_rgbt = torch.split(score_map_ctr, 1, dim=0)
+        #print("score_map_ctr_rgbt shape  ", score_map_ctr_rgbt[0].shape)
+        # assert gt_score_map is None
+        if gt_score_map is None:
+            bbox = self.cal_bbox(score_map_ctr, size_map, offset_map)
+        else:
+            bbox = self.cal_bbox(gt_score_map.unsqueeze(1), size_map, offset_map)
+        
+
+
+        return score_map_ctr, bbox, size_map, offset_map
+
+    def cal_bbox(self, score_map_ctr_rgbt, size_map_rgbt, offset_map_rgbt, return_score=False):
+        max_score1, idx1 = (torch.max(score_map_ctr_rgbt[0].flatten(1), dim=1, keepdim=True))                                        
+        max_score2, idx2 = (torch.max(score_map_ctr_rgbt[1].flatten(1), dim=1, keepdim=True))
+        
+
+        sm = torch.split(size_map_rgbt, 1, dim=0)
+        om = torch.split(offset_map_rgbt, 1, dim=0)
+
+        if max_score1 >= max_score2:
+            max_score = max_score1
+            idx = idx1
+            size_map = sm[0]
+            offset_map = om[0]
+        else:
+            max_score = max_score2
+            idx = idx2
+            size_map = sm[1]
+            offset_map = om[1]
+            #size_map = size_map_rgbt[1]
+            #offset_map = offset_map_rgbt[1]
+
+        
+
+        idx_y = idx // self.feat_sz
+        idx_x = idx % self.feat_sz
+        
+        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
+        
+        size = size_map.flatten(2).gather(dim=2, index=idx)
+        
+        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+
+
+        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
+        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
+        # cx, cy, w, h
+        bbox = torch.cat([(idx_x.to(torch.float) + offset[:, :1]) / self.feat_sz,
+                          (idx_y.to(torch.float) + offset[:, 1:]) / self.feat_sz,
+                          size.squeeze(-1)], dim=1)
+        
+        if return_score:
+            return bbox, max_score
+        return bbox
+
+    def get_pred(self, score_map_ctr, size_map, offset_map):
+        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
+        idx_y = idx // self.feat_sz
+        idx_x = idx % self.feat_sz
+
+        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
+        size = size_map.flatten(2).gather(dim=2, index=idx)
+        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+
+        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
+        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
+        return size * self.feat_sz, offset
+
+    def get_score_map(self, x):
+
+        
+        def _sigmoid(x):
+            y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4)
+            return y
+
+        # ctr branch
+        x_ctr1 = self.conv1_ctr(x)
+        x_ctr2 = self.conv2_ctr(x_ctr1)
+        x_ctr3 = self.conv3_ctr(x_ctr2)
+        x_ctr4 = self.conv4_ctr(x_ctr3)
+        score_map_ctr = self.conv5_ctr(x_ctr4)
+
+        # offset branch
+        x_offset1 = self.conv1_offset(x)
+        x_offset2 = self.conv2_offset(x_offset1)
+        x_offset3 = self.conv3_offset(x_offset2)
+        x_offset4 = self.conv4_offset(x_offset3)
+        score_map_offset = self.conv5_offset(x_offset4)
+
+        # size branch
+        x_size1 = self.conv1_size(x)
+        x_size2 = self.conv2_size(x_size1)
+        x_size3 = self.conv3_size(x_size2)
+        x_size4 = self.conv4_size(x_size3)
+        score_map_size = self.conv5_size(x_size4)
+        return _sigmoid(score_map_ctr), _sigmoid(score_map_size), score_map_offset
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, BN=False):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        if BN:
+            self.layers = nn.ModuleList(nn.Sequential(nn.Linear(n, k), nn.BatchNorm1d(k))
+                                        for n, k in zip([input_dim] + h, h + [output_dim]))
+        else:
+            self.layers = nn.ModuleList(nn.Linear(n, k)
+                                        for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def build_box_head(cfg, hidden_dim):
+    stride = cfg.MODEL.BACKBONE.STRIDE
+
+    if cfg.MODEL.HEAD.TYPE == "MLP":
+        mlp_head = MLP(hidden_dim, hidden_dim, 4, 3)  # dim_in, dim_hidden, dim_out, 3 layers
+        return mlp_head
+    elif "CORNER" in cfg.MODEL.HEAD.TYPE:
+        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
+        channel = getattr(cfg.MODEL, "NUM_CHANNELS", 256)
+        print("head channel: %d" % channel)
+        if cfg.MODEL.HEAD.TYPE == "CORNER":
+            corner_head = Corner_Predictor(inplanes=cfg.MODEL.HIDDEN_DIM, channel=channel,
+                                           feat_sz=feat_sz, stride=stride)
+        else:
+            raise ValueError()
+        return corner_head
+    elif cfg.MODEL.HEAD.TYPE == "CENTER":
+        in_channel = hidden_dim
+        out_channel = cfg.MODEL.HEAD.NUM_CHANNELS
+        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
+        center_head = CenterPredictor(inplanes=in_channel, channel=out_channel,
+                                      feat_sz=feat_sz, stride=stride)
+        return center_head
+    else:
+        raise ValueError("HEAD TYPE %s is not supported." % cfg.MODEL.HEAD_TYPE)

lib/models/layers/patch_embed.py

@@ -0,0 +1,32 @@
+import torch.nn as nn
+
+from timm.models.layers import to_2tuple
+
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        # allow different input size
+        # B, C, H, W = x.shape
+        # _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
+        # _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x

lib/models/layers/position_encoding.py

@@ -0,0 +1,154 @@
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+from lib.utils.misc import NestedTensor
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        mask = tensor_list.mask
+        assert mask is not None
+        not_mask = ~mask # (b,h,w)
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)  # cumulative sum along axis 1 (h axis) --> (b, h, w)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)  # cumulative sum along axis 2 (w axis) --> (b, h, w)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale  # 2pi * (y / sigma(y))
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale  # 2pi * (x / sigma(x))
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)  # (0,1,2,...,d/2)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t # (b,h,w,d/2)
+        pos_y = y_embed[:, :, :, None] / dim_t # (b,h,w,d/2)
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) # (b,h,w,d/2)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) # (b,h,w,d/2)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)  # (b,h,w,d)
+        return pos
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.row_embed = nn.Embedding(50, num_pos_feats)
+        self.col_embed = nn.Embedding(50, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        h, w = x.shape[-2:]
+        i = torch.arange(w, device=x.device)
+        j = torch.arange(h, device=x.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
+        return pos  # (H,W,C) --> (C,H,W) --> (1,C,H,W) --> (B,C,H,W)
+
+
+class PositionEmbeddingNone(nn.Module):
+    """
+    No positional encoding.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.n_dim = num_pos_feats * 2
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        b, _, h, w = x.size()
+        return torch.zeros((b, self.n_dim, h, w), device=x.device)  # (B, C, H, W)
+
+
+def build_position_encoding(cfg):
+    N_steps = cfg.MODEL.HIDDEN_DIM // 2
+    if cfg.MODEL.POSITION_EMBEDDING in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    elif cfg.MODEL.POSITION_EMBEDDING in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(N_steps)
+    elif cfg.MODEL.POSITION_EMBEDDING in ('None', ):
+        print("Not using positional encoding.")
+        position_embedding = PositionEmbeddingNone(N_steps)
+    else:
+        raise ValueError(f"not supported {cfg.MODEL.POSITION_EMBEDDING}")
+
+    return position_embedding
+
+
+class PositionEmbeddingLearned_new(nn.Module):
+    """
+    Absolute pos embedding, learned. (allow users to specify the size)
+    """
+    def __init__(self, num_pos_feats=256, sz=20):
+        super().__init__()
+        self.sz = sz
+        self.row_embed = nn.Embedding(sz, num_pos_feats)
+        self.col_embed = nn.Embedding(sz, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, bs):
+        """bs: batch size"""
+        h, w = self.sz, self.sz
+        i = torch.arange(w, device=self.col_embed.weight.device)
+        j = torch.arange(h, device=self.row_embed.weight.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        # pos = torch.cat([
+        #     x_emb.unsqueeze(0).repeat(h, 1, 1),
+        #     y_emb.unsqueeze(1).repeat(1, w, 1),
+        # ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(bs, 1, 1, 1)
+
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).repeat(1, bs, 1)
+        # print(bs, pos)
+        return pos  # (H,W,C) --> (C,H,W) --> (1,C,H,W) --> (B,C,H,W)
+
+
+# def build_position_encoding_new(cfg, sz):
+#     N_steps = cfg.MODEL.HIDDEN_DIM // 2
+#     position_embedding = PositionEmbeddingLearned_new(N_steps, sz)
+#     return position_embedding
+
+
+
+def build_position_encoding(embed_dim, sz=256):
+    N_steps = embed_dim // 2
+    position_embedding = PositionEmbeddingLearned_new(N_steps, sz)
+    return position_embedding
+

lib/models/layers/rpe.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_
+
+
+def generate_2d_relative_positional_encoding_index(z_shape, x_shape):
+    '''
+        z_shape: (z_h, z_w)
+        x_shape: (x_h, x_w)
+    '''
+    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
+    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))
+
+    z_2d_index_h = z_2d_index_h.flatten(0)
+    z_2d_index_w = z_2d_index_w.flatten(0)
+    x_2d_index_h = x_2d_index_h.flatten(0)
+    x_2d_index_w = x_2d_index_w.flatten(0)
+
+    diff_h = z_2d_index_h[:, None] - x_2d_index_h[None, :]
+    diff_w = z_2d_index_w[:, None] - x_2d_index_w[None, :]
+
+    diff = torch.stack((diff_h, diff_w), dim=-1)
+    _, indices = torch.unique(diff.view(-1, 2), return_inverse=True, dim=0)
+    return indices.view(z_shape[0] * z_shape[1], x_shape[0] * x_shape[1])
+
+
+def generate_2d_concatenated_self_attention_relative_positional_encoding_index(z_shape, x_shape):
+    '''
+        z_shape: (z_h, z_w)
+        x_shape: (x_h, x_w)
+    '''
+    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
+    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))
+
+    z_2d_index_h = z_2d_index_h.flatten(0)
+    z_2d_index_w = z_2d_index_w.flatten(0)
+    x_2d_index_h = x_2d_index_h.flatten(0)
+    x_2d_index_w = x_2d_index_w.flatten(0)
+
+    concatenated_2d_index_h = torch.cat((z_2d_index_h, x_2d_index_h))
+    concatenated_2d_index_w = torch.cat((z_2d_index_w, x_2d_index_w))
+
+    diff_h = concatenated_2d_index_h[:, None] - concatenated_2d_index_h[None, :]
+    diff_w = concatenated_2d_index_w[:, None] - concatenated_2d_index_w[None, :]
+
+    z_len = z_shape[0] * z_shape[1]
+    x_len = x_shape[0] * x_shape[1]
+    a = torch.empty((z_len + x_len), dtype=torch.int64)
+    a[:z_len] = 0
+    a[z_len:] = 1
+    b=a[:, None].repeat(1, z_len + x_len)
+    c=a[None, :].repeat(z_len + x_len, 1)
+
+    diff = torch.stack((diff_h, diff_w, b, c), dim=-1)
+    _, indices = torch.unique(diff.view((z_len + x_len) * (z_len + x_len), 4), return_inverse=True, dim=0)
+    return indices.view((z_len + x_len), (z_len + x_len))
+
+
+def generate_2d_concatenated_cross_attention_relative_positional_encoding_index(z_shape, x_shape):
+    '''
+        z_shape: (z_h, z_w)
+        x_shape: (x_h, x_w)
+    '''
+    z_2d_index_h, z_2d_index_w = torch.meshgrid(torch.arange(z_shape[0]), torch.arange(z_shape[1]))
+    x_2d_index_h, x_2d_index_w = torch.meshgrid(torch.arange(x_shape[0]), torch.arange(x_shape[1]))
+
+    z_2d_index_h = z_2d_index_h.flatten(0)
+    z_2d_index_w = z_2d_index_w.flatten(0)
+    x_2d_index_h = x_2d_index_h.flatten(0)
+    x_2d_index_w = x_2d_index_w.flatten(0)
+
+    concatenated_2d_index_h = torch.cat((z_2d_index_h, x_2d_index_h))
+    concatenated_2d_index_w = torch.cat((z_2d_index_w, x_2d_index_w))
+
+    diff_h = x_2d_index_h[:, None] - concatenated_2d_index_h[None, :]
+    diff_w = x_2d_index_w[:, None] - concatenated_2d_index_w[None, :]
+
+    z_len = z_shape[0] * z_shape[1]
+    x_len = x_shape[0] * x_shape[1]
+
+    a = torch.empty(z_len + x_len, dtype=torch.int64)
+    a[: z_len] = 0
+    a[z_len:] = 1
+    c = a[None, :].repeat(x_len, 1)
+
+    diff = torch.stack((diff_h, diff_w, c), dim=-1)
+    _, indices = torch.unique(diff.view(x_len * (z_len + x_len), 3), return_inverse=True, dim=0)
+    return indices.view(x_len, (z_len + x_len))
+
+
+class RelativePosition2DEncoder(nn.Module):
+    def __init__(self, num_heads, embed_size):
+        super(RelativePosition2DEncoder, self).__init__()
+        self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads, embed_size)))
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+
+    def forward(self, attn_rpe_index):
+        '''
+            Args:
+                attn_rpe_index (torch.Tensor): (*), any shape containing indices, max(attn_rpe_index) < embed_size
+            Returns:
+                torch.Tensor: (1, num_heads, *)
+        '''
+        return self.relative_position_bias_table[:, attn_rpe_index].unsqueeze(0)

lib/models/layers/transformer_dec.py

@@ -0,0 +1,415 @@
+import copy
+from typing import Optional, List
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+import time
+
+
+def check_inf(tensor):
+    return torch.isinf(tensor.detach()).any()
+
+
+def check_nan(tensor):
+    return torch.isnan(tensor.detach()).any()
+
+
+def check_valid(tensor, type_name):
+    if check_inf(tensor):
+        print("%s is inf." % type_name)
+    if check_nan(tensor):
+        print("%s is nan" % type_name)
+
+
+def generate_square_subsequent_mask(sz):
+    r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+        Unmasked positions are filled with float(0.0).
+    """
+
+    # each token only can see tokens before them
+    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+    mask = mask.float().masked_fill(mask == 0, float(
+        '-inf')).masked_fill(mask == 1, float(0.0))
+    return mask
+
+
+class Transformer_dec(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=768, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False, divide_norm=False):
+        super().__init__()
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before, divide_norm=divide_norm)
+        decoder_norm = nn.LayerNorm(d_model)
+
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
+                                          return_intermediate=return_intermediate_dec)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+        self.d_feed = dim_feedforward
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self,tgt, feat, pos, query_embed,mask):
+        memory = feat  # output of hivit
+        hs = self.decoder(tgt, memory, pos=pos, query_pos=query_embed,memory_key_padding_mask=mask)
+        return hs
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+
+        intermediate = []
+        for layer_id, layer in enumerate(self.layers):
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output[0] = self.norm(output[0])
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+        return output
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False, divide_norm=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout,batch_first=True)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout,batch_first=True)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+        self.divide_norm = divide_norm
+        self.scale_factor = float(d_model // nhead) ** 0.5
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt_all, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        # print(tgt.shape, memory.shape)
+        # print(len(tgt_all))
+        q_initial = tgt_all[0]
+        kv_initial = tgt_all[1]
+        q = kv_initial
+        k = tgt = kv_initial  # self.with_pos_embed(tgt, query_pos)
+        # print(q.shape, k.shape, tgt.shape, 'tgt kv')
+        tgt2 = self.self_attn(self.with_pos_embed(q,query_pos), self.with_pos_embed(k,query_pos), value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=memory_key_padding_mask)[0]
+        # print(tgt.shape, tgt2.shape, 'tgt 1')
+        tgt = q + self.dropout1(tgt2)
+        # print(tgt.shape, tgt2.shape, 'tgt 2')
+        kv_update = self.norm1(tgt)
+
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(q_initial, pos),
+                                   key=self.with_pos_embed(kv_update, query_pos),
+                                   value=kv_update, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = q_initial + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+
+        # tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        # tgt = tgt + self.dropout3(tgt2)
+        # tgt = self.norm3(tgt)
+        # print(tgt.shape, 'tgt 2')
+        return [tgt, kv_update]
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos,query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos,query_pos)
+
+
+class Transformer_dec_with_mask(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=768, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False, divide_norm=False):
+        super().__init__()
+        decoder_layer = TransformerDecoderWithMaskLayer(d_model, nhead, dim_feedforward,
+                                                        dropout, activation, normalize_before, divide_norm=divide_norm)
+        decoder_norm = nn.LayerNorm(d_model)
+
+        self.decoder = TransformerDecoderWithMask(decoder_layer, num_decoder_layers, decoder_norm,
+                                                  return_intermediate=return_intermediate_dec)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+        self.d_feed = dim_feedforward
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, q, kv, feat_mask, q_pos, kv_pos=None):
+        memory = kv  # output of hivit
+        hs = self.decoder(q, memory, feat_mask, pos=kv_pos, query_pos=q_pos)
+        return hs
+
+
+class TransformerDecoderWithMask(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory, feat_mask,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None,
+                pre_query_pos: Optional[Tensor] = None):
+        output = tgt
+
+        intermediate = []
+        for layer_id, layer in enumerate(self.layers):
+            output = layer(output, memory, feat_mask,
+                           tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos, pre_query_pos=pre_query_pos)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+        return output
+
+
+class TransformerDecoderWithMaskLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False, divide_norm=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+        self.divide_norm = divide_norm
+        self.scale_factor = float(d_model // nhead) ** 0.5
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt_all, memory, feat_mask,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None,
+                     pre_query_pos: Optional[Tensor] = None):
+        q = tgt_all
+        k = memory * feat_mask
+        tgt2 = self.self_attn(self.with_pos_embed(q, query_pos ), self.with_pos_embed(k, pos), value=k,
+                              attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = q + self.dropout1(tgt2)
+        tgt = self.norm2(tgt)
+
+        # memory_1 = self.norm1(tgt)
+        #
+        # tgt = tgt_all
+        # tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+        #                            key=self.with_pos_embed(memory_1, pos),
+        #                            value=memory_1, attn_mask=memory_mask,
+        #                            key_padding_mask=memory_key_padding_mask)[0]
+        # tgt = tgt + self.dropout2(tgt2)
+        # tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+
+        return tgt
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory, feat_mask,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None,
+                pre_query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos)
+        return self.forward_post(tgt, memory, feat_mask, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos, pre_query_pos)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_transformer_dec(cfg, hidden_dim):
+    return Transformer_dec(
+        d_model=hidden_dim,
+        dropout=0.1,
+        nhead=8,
+        dim_feedforward=hidden_dim,
+        num_encoder_layers=0,
+        num_decoder_layers=2,
+        normalize_before=False,
+        return_intermediate_dec=False,
+        divide_norm=False
+    )
+
+
+def build_transformer_dec_with_mask(cfg, hidden_dim):
+    return Transformer_dec_with_mask(
+        d_model=hidden_dim,
+        dropout=0.1,
+        nhead=8,
+        dim_feedforward=hidden_dim,
+        num_encoder_layers=0,
+        num_decoder_layers=2,
+        normalize_before=False,
+        return_intermediate_dec=False,
+        divide_norm=False
+    )
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

lib/models/transformers/__init__.py

@@ -0,0 +1,2 @@
+from .transformer import build_decoder, VisionLanguageFusionModule,build_temporal_decoder,build_text_prompt_decoder
+from .position_encoding import PositionEmbeddingSine1D

lib/models/transformers/position_encoding.py

@@ -0,0 +1,154 @@
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+from lib.utils.misc import NestedTensor
+
+
+class PositionEmbeddingSine1D(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=256, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors # [B, L, C]
+        mask = tensor_list.mask # [B, L]
+        assert mask is not None
+        not_mask = ~mask
+        x_embed = not_mask.cumsum(1, dtype=torch.float32)  # [B, L]
+        if self.normalize:
+            eps = 1e-6
+            x_embed = x_embed / (x_embed[:, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, None] / dim_t  # [B, L, C]
+        pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)  # [B, L, C]
+        # pos = pos_x.permute(0, 2, 1)    # [B, C, L]
+        return pos_x
+
+
+class PositionEmbeddingSine2D(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, mask):
+        # mask: (B, H, W)
+        not_mask = ~mask  # (batch,h,w)  ~mask:按位取反
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)  # cumulative sum along axis 1 (h axis) --> (b, h, w)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)  # cumulative sum along axis 2 (w axis) --> (b, h, w)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale  # 2π * (y / sigma(y))  (16,8,8)
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale  # 2π * (x / sigma(x))  (16,8,8)
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=mask.device)  # (0,1,2,...,d/2)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t  # (b,h,w,d/2)
+        pos_y = y_embed[:, :, :, None] / dim_t  # (b,h,w,d/2)
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)  # (b,h,w,d/2) (16,8,8,128)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)  # (b,h,w,d/2) (16,8,8,128)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)  # (b,h,w,d) (16,8,8,256) to (16,256,8,8)
+        return pos
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.embed = nn.Embedding(in_dim, out_dim)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.embed.weight)
+
+    def forward(self, x):
+        # x: (B, 5, C) or (B, mask, C) or (B, bbox+mask, C)
+        n = x.size(1)
+        i = torch.arange(n, device=x.device)
+        pos = self.embed(i).unsqueeze(0).repeat(x.size(0), 1, 1) # (N,C) --> (1,N,C) --> (B,N,C)
+        return pos
+
+
+class PositionEmbeddingNone(nn.Module):
+    """
+    No positional encoding.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.n_dim = num_pos_feats * 2
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        b, _, h, w = x.size()
+        return torch.zeros((b, self.n_dim, h, w), device=x.device)  # (B, C, H, W)
+
+
+def build_memory_position_encoding(cfg):
+    N_steps = cfg.MODEL.DECODER.HIDDEN_DIM // 2  # N_steps = 128
+    if cfg.MODEL.DECODER.MEMORY_POSITION_EMBEDDING in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine2D(N_steps, normalize=True)
+    elif cfg.MODEL.DECODER.MEMORY_POSITION_EMBEDDING in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(N_steps)
+    elif cfg.MODEL.DECODER.MEMORY_POSITION_EMBEDDING in ('None', ):
+        print("Not using positional encoding.")
+        position_embedding = PositionEmbeddingNone(N_steps)
+    else:
+        raise ValueError(f"not supported {cfg.MODEL.DECODER.MEMORY_POSITION_EMBEDDING}")
+
+    return position_embedding
+
+
+def build_query_position_encoding(cfg):
+    bbox_task = cfg.TRAIN.BBOX_TASK
+    language_task = getattr(cfg.TRAIN, "LANGUAGE_TASK", False)
+        
+    if bbox_task and language_task:      # bbox and language tasks
+        seq_in_dim = (4+1) + (1+1)
+    elif bbox_task and not language_task:  # bbox level
+        seq_in_dim = 4+1
+
+    N_steps = cfg.MODEL.DECODER.HIDDEN_DIM
+    if cfg.MODEL.DECODER.QUERY_POSITION_EMBEDDING in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine2D(N_steps, normalize=True)
+    elif cfg.MODEL.DECODER.QUERY_POSITION_EMBEDDING in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(in_dim=seq_in_dim, out_dim=N_steps)
+    elif cfg.MODEL.DECODER.QUERY_POSITION_EMBEDDING in ('None', ):
+        print("Not using positional encoding.")
+        position_embedding = PositionEmbeddingNone(N_steps)
+    else:
+        raise ValueError(f"not supported {cfg.MODEL.DECODER.QUERY_POSITION_EMBEDDING}")
+
+    return position_embedding

lib/models/transformers/transformer.py

@@ -0,0 +1,498 @@
+import copy
+from tkinter import N
+from typing import Optional, List
+import warnings
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+import math
+from functools import partial
+from timm.models.helpers import build_model_with_cfg, named_apply, adapt_input_conv
+from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
+from .position_encoding import build_memory_position_encoding, build_query_position_encoding
+
+from lib.models.transformers.position_encoding import PositionEmbeddingNone,PositionEmbeddingLearned
+
+
+def with_pos_embed(tensor, pos: Optional[Tensor]):
+    return tensor if pos is None else tensor + pos
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+
+
+    def forward(self, query, key, value=None,
+                attn_mask=None, key_padding_mask=None,
+                need_weights=False):
+        """
+        query: [B, N, C]
+        attn_mask: [N, N] torch.float32
+        key_padding_mask: [B, N] torch.bool
+        """
+        B, q_N, C = query.shape
+        
+        if attn_mask is not None:
+            assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
+                attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
+                'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
+            if attn_mask.dtype == torch.uint8:
+                warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+                attn_mask = attn_mask.to(torch.bool)
+
+            if attn_mask.dim() == 2:
+                attn_mask = attn_mask.unsqueeze(0)
+                if list(attn_mask.size()) != [1, query.size(1), key.size(1)]:
+                    raise RuntimeError('The size of the 2D attn_mask is not correct.')
+            elif attn_mask.dim() == 3:
+                if list(attn_mask.size()) != [B * self.num_heads, query.size(1), key.size(1)]:
+                    raise RuntimeError('The size of the 3D attn_mask is not correct.')
+            else:
+                raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
+            # attn_mask's dim is 3 now.
+
+        # convert key_padding_mask to bool
+        if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+            warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+            key_padding_mask = key_padding_mask.to(torch.bool)
+
+        q = self.q(query).reshape(B, q_N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)  # B,head,HW,C/head
+        k = self.k(key).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        v = self.v(value).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale # (B, head, N, N)
+        
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn.masked_fill_(attn_mask, float('-inf'))
+            else:
+                attn += attn_mask
+        
+        if key_padding_mask is not None:
+            attn = attn.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2),
+                float('-inf'),
+            )
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, q_N, C)
+        x = self.proj(x)
+
+        if need_weights:
+            # average attention weights over heads
+            return x, attn.sum(dim=1) / self.num_heads  # attn: (B, q_len, k_len)
+        else:
+            return x, None
+
+
+class MLP(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): # nn.ReLU
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.drop = nn.Dropout(drop)
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(self, d_model=256, nhead=8, num_decoder_layers=0, mlp_ratio=1, 
+                dropout=0., vocab_size=1001, memory_pos_embedding=None, query_pos_embedding=None):
+        super().__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+        self.query_embedding = nn.Embedding(vocab_size, self.d_model)  # (1001, 256)
+        self.memory_pos_embedding = memory_pos_embedding
+        self.query_pos_embedding = query_pos_embedding
+
+        decoder_layer = TransformerDecoderLayer(d_model=d_model, num_heads=nhead, mlp_ratio=mlp_ratio,
+                                                attn_drop=dropout, proj_drop=dropout)
+        if num_decoder_layers == 0:
+            self.layers = None
+        else:
+            self.layers = _get_clones(decoder_layer, num_decoder_layers)
+
+        self.num_decoder_layers = num_decoder_layers
+        self.scale_factor = float(d_model // nhead) ** 0.5
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out = fan_out // m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+
+    def memory_mask_pos_enc(self, attn_mask, feat_sz):
+        """
+        attn_mask: (B, img_H, img_W)
+        feat_sz: feature size
+        """
+        batch_size = attn_mask.size(0)
+        attn_mask = attn_mask.to(torch.float32)
+        attn_mask = F.interpolate(attn_mask.unsqueeze(1), size=(feat_sz, feat_sz)).to(torch.bool).squeeze(1)
+
+        pos_embeds = self.memory_pos_embedding(attn_mask) # sine position encoding  (B, C, feat_sz, feat_sz)
+
+        attn_mask = attn_mask.view(batch_size, -1)
+        pos_embeds = pos_embeds.view(batch_size, self.d_model, -1).transpose(1, 2)
+
+        return attn_mask, pos_embeds
+
+
+    def tri_mask(self, length):
+        mask = (torch.triu(torch.ones(length, length)) == 1).float().transpose(0, 1)
+        mask.masked_fill_(mask == 0, float('-inf'))
+        mask.masked_fill_(mask == 1, float(0.))
+        return mask
+
+
+    def forward(self, query, memory, memory_key_padding_mask, memory_pos,
+                return_intermediate_output=False, is_inference=False):
+        """
+        query:        (B, 5, C) or (B, 18, C) or (B, bbox+mask, C)
+        memory:       (B, HW, C)
+        memory_key_padding_mask:  (B, HW)
+        memory_pos:   (B, HW, C)
+        """
+        # learnable pos encoding
+        query_pos = self.query_pos_embedding(query)
+
+        query_mask = self.tri_mask(query.size(1)).to(query.device)
+
+        aux_logits = []
+        output = query
+        for i in range(len(self.layers)):
+            output, attn_weights = self.layers[i](output, memory,
+                                                query_pos=query_pos, memory_pos=memory_pos,
+                                                query_mask=query_mask,
+                                                memory_key_padding_mask=memory_key_padding_mask,
+                                                need_weights=False)
+            if return_intermediate_output and i < len(self.layers)-1:
+                aux_logits.append(output)
+        
+        if is_inference:
+            return output
+        else:
+            return output, aux_logits
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(self, d_model=256, num_heads=8, mlp_ratio=8,
+                 attn_drop=0.1, proj_drop=0.1):
+        super().__init__()
+        # self_attn
+        # self.self_attn = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
+        # self.dropout1 = nn.Dropout(proj_drop)
+        # self.norm1 = nn.LayerNorm(d_model)
+        #
+        # self.self_attn_det = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
+        # self.dropout1_det = nn.Dropout(proj_drop)
+        # self.norm1_det = nn.LayerNorm(d_model)
+
+        # cross_attn
+        self.cross_attn = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
+        self.dropout2 = nn.Dropout(proj_drop)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.cross_attn_det = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
+        self.dropout2_det = nn.Dropout(proj_drop)
+        self.norm2_det = nn.LayerNorm(d_model)
+
+        # mlp
+        dim_feedforward = d_model * mlp_ratio
+        self.MLP = MLP(in_features=d_model, hidden_features=dim_feedforward, out_features=d_model, drop=proj_drop)
+        self.dropout3 = nn.Dropout(proj_drop)
+        self.norm3 = nn.LayerNorm(d_model)
+
+        self.MLP_det = MLP(in_features=d_model, hidden_features=dim_feedforward, out_features=d_model, drop=proj_drop)
+        self.dropout3_det = nn.Dropout(proj_drop)
+        self.norm3_det = nn.LayerNorm(d_model)
+
+        # final
+        self.norm_final = nn.LayerNorm(d_model)
+
+
+    def forward(self, query, memory,
+                query_pos: Optional[Tensor] = None,
+                memory_pos: Optional[Tensor] = None,
+                query_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                need_weights: bool = False):
+        """
+        query: (B, 5, C)
+        memory: (B, HW, C)
+        query_pos: (B, 5, C)
+        memory_pos: (B, HW, C)
+        query_mask: (5,5)
+        memory_key_padding_mask: (B, HW)
+        """
+        # query self attn
+        # q = k = with_pos_embed(query, query_pos)
+        # query2, selfattn_weights = self.self_attn(q, k, value=query, attn_mask=query_mask, need_weights=need_weights)
+        # query = query + self.dropout1(query2)
+        # query = self.norm1(query)
+        #
+        # # memory self attn
+        # q = k = with_pos_embed(memory, memory_pos)
+        # query2, selfattn_weights = self.self_attn_det(q, k, value=memory, attn_mask=memory_key_padding_mask, need_weights=need_weights)
+        # memory = memory + self.dropout1_det(query2)
+        # memory = self.norm1_det(memory)
+
+        # query to memory
+        query2, crossattn_weights = self.cross_attn(query=with_pos_embed(query, query_pos),
+                                                    key=with_pos_embed(memory, memory_pos),
+                                                    value=memory,
+                                                    key_padding_mask=memory_key_padding_mask,
+                                                    need_weights=need_weights)
+        query = query + self.dropout2(query2)
+        query = self.norm2(query)
+
+        #  memory to query
+        query2, crossattn_weights = self.cross_attn_det(query=with_pos_embed(memory, memory_pos),
+                                                    key=with_pos_embed(query, query_pos),
+                                                    value=query,
+                                                    key_padding_mask=memory_key_padding_mask,
+                                                    need_weights=need_weights)
+        memory = memory + self.dropout2_det(query2)
+        memory = self.norm2_det(memory)
+
+
+        # mlp
+        query2 = self.MLP(query)
+        query = query + self.dropout3(query2)
+        query = self.norm3(query)
+
+        memory = self.MLP_det(memory)
+        memory = memory + self.dropout3_det(memory)
+        memory = self.norm3_det(memory)
+
+        ### final fusion
+        final_x = query[:,4:] + memory[:,4:]
+        final = torch.cat([query[:,:4], memory[:,:4] , final_x ],dim=1)
+        query = self.norm_final(final)
+
+        if need_weights:
+            return query, crossattn_weights
+        else:
+            return query, None
+class TransformerDecoder_v0(nn.Module):
+    def __init__(self, d_model=256, nhead=8, num_decoder_layers=0, mlp_ratio=1,
+                 dropout=0.):
+        super().__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+
+        decoder_layer = TransformerDecoderLayer(d_model=d_model, num_heads=nhead, mlp_ratio=mlp_ratio,
+                                                attn_drop=dropout, proj_drop=dropout)
+        if num_decoder_layers == 0:
+            self.layers = None
+        else:
+            self.layers = _get_clones(decoder_layer, num_decoder_layers)
+
+        self.num_decoder_layers = num_decoder_layers
+        self.scale_factor = float(d_model // nhead) ** 0.5
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out = fan_out // m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def memory_mask_pos_enc(self, attn_mask, feat_sz):
+        """
+        attn_mask: (B, img_H, img_W)
+        feat_sz: feature size
+        """
+        batch_size = attn_mask.size(0)
+        attn_mask = attn_mask.to(torch.float32)
+        attn_mask = F.interpolate(attn_mask.unsqueeze(1), size=(feat_sz, feat_sz)).to(torch.bool).squeeze(1)
+
+        pos_embeds = self.memory_pos_embedding(attn_mask)  # sine position encoding  (B, C, feat_sz, feat_sz)
+
+        attn_mask = attn_mask.view(batch_size, -1)
+        pos_embeds = pos_embeds.view(batch_size, self.d_model, -1).transpose(1, 2)
+
+        return attn_mask, pos_embeds
+
+
+    def forward(self, query, memory, query_pos, memory_pos):
+        """
+        query:        (B, 5, C) or (B, 18, C) or (B, bbox+mask, C)
+        memory:       (B, HW, C)
+        memory_key_padding_mask:  (B, HW)
+        memory_pos:   (B, HW, C)
+        """
+        # learnable pos encoding
+        # query_pos = self.query_pos_embedding(query)
+
+        output = query
+        for i in range(len(self.layers)):
+            output, attn_weights = self.layers[i](output, memory,
+                                                  query_pos=query_pos, memory_pos=memory_pos,
+                                                  memory_key_padding_mask=None,
+                                                  need_weights=False)
+
+
+        return output
+
+class VisionLanguageFusionModule(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., num_vlfusion_layers=0,
+                vl_input_type='separate'):
+        super().__init__()
+        # self.multihead_attn = MultiheadAttention(dim, num_heads=num_heads, 
+        #                                     qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop)
+        
+        self.vl_input_type = vl_input_type
+        VLFusion_layer = MultiheadAttention(dim, num_heads=num_heads, 
+                                            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop)
+        if num_vlfusion_layers == 0:
+            self.VLFusion_layers = None
+        else:
+            self.VLFusion_layers = _get_clones(VLFusion_layer, num_vlfusion_layers)
+
+    def forward(self, query, memory,
+                query_pos: Optional[Tensor] = None,
+                memory_pos: Optional[Tensor] = None,
+                query_mask: Optional[Tensor] = None,
+                query_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                need_weights: bool = False):
+
+        if self.vl_input_type == 'separate':
+            output = query
+        elif self.vl_input_type == 'concat':
+            output, query_pos = torch.cat([memory, query], dim=1), torch.cat([memory_pos, query_pos], dim=1)
+            memory, memory_pos = output.clone(), query_pos.clone()
+            memory_key_padding_mask = torch.cat([memory_key_padding_mask, query_key_padding_mask], dim=1)
+
+        for layer in self.VLFusion_layers:
+            output, attn_weights = layer(query=with_pos_embed(output, query_pos),
+                                        key=with_pos_embed(memory, memory_pos),
+                                        value=memory,
+                                        key_padding_mask=memory_key_padding_mask,
+                                        need_weights=need_weights)      # attn_weights: (B, q_len, k_len)
+            
+            output = query * output         # todo: 残乘，应该是和前一个block的输出相乘吧，而不是始终和初始的输入相乘？ 或者把这个放在 for循环外边
+
+
+        return output
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_decoder(cfg):
+    memory_pos_embedding = build_memory_position_encoding(cfg)
+    query_pos_embedding = build_query_position_encoding(cfg)
+    return TransformerDecoder(
+        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+        dropout=cfg.MODEL.DECODER.DROPOUT,
+        nhead=cfg.MODEL.DECODER.NUM_HEADS,
+        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
+        num_decoder_layers=cfg.MODEL.DECODER.DEC_LAYERS,
+        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
+        memory_pos_embedding=memory_pos_embedding,
+        query_pos_embedding=query_pos_embedding,
+    )
+
+from lib.models.transformers.position_encoding import PositionEmbeddingNone,PositionEmbeddingLearned
+def build_temporal_decoder(cfg):
+    memory_pos_embedding = PositionEmbeddingLearned(15,256)
+    query_pos_embedding = PositionEmbeddingLearned(15,256)
+    return TransformerDecoder_v0(
+        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+        dropout=cfg.MODEL.DECODER.DROPOUT,
+        nhead=cfg.MODEL.DECODER.NUM_HEADS,
+        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
+        num_decoder_layers=2,
+        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
+        memory_pos_embedding=memory_pos_embedding,
+        query_pos_embedding=query_pos_embedding,
+    )
+
+def build_text_prompt_decoder(cfg):
+    memory_pos_embedding = PositionEmbeddingLearned(15,512)
+    query_pos_embedding = PositionEmbeddingLearned(15,512)
+    return TransformerDecoder_v0(
+        d_model=512,
+        dropout=cfg.MODEL.DECODER.DROPOUT,
+        nhead=cfg.MODEL.DECODER.NUM_HEADS,
+        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
+        num_decoder_layers=2,
+        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
+        memory_pos_embedding=memory_pos_embedding,
+        query_pos_embedding=query_pos_embedding,
+    )
+
+def build_img_prompt_decoder(cfg):
+    memory_pos_embedding = PositionEmbeddingLearned(15,256)
+    query_pos_embedding = PositionEmbeddingLearned(15,256)
+    return TransformerDecoder_v0(
+        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+        dropout=cfg.MODEL.DECODER.DROPOUT,
+        nhead=8,
+        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
+        num_decoder_layers=2,
+        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
+        memory_pos_embedding=memory_pos_embedding,
+        query_pos_embedding=query_pos_embedding,
+    )
+
+
+def build_rgb_det_decoder(cfg):
+    return TransformerDecoder_v0(
+        d_model=512,
+        dropout=0.1,
+        nhead=8,
+        mlp_ratio=1,
+        num_decoder_layers=1
+    )
+
+
+
+

lib/models/transformers/transformer_visual.py

@@ -0,0 +1,385 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+DETR Transformer class.
+
+Copy-paste from torch.nn.Transformer with modifications:
+    * positional encodings are passed in MHattention
+    * extra LN at the end of encoder is removed
+    * decoder returns a stack of activations from all decoding layers
+"""
+import copy
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+from lib.models.utils import QuickGELU
+
+
+class Transformer(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False):
+        super().__init__()
+
+        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        if num_decoder_layers > 0:
+            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                    dropout, activation, normalize_before)
+            decoder_norm = nn.LayerNorm(d_model)
+            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
+                                              return_intermediate=return_intermediate_dec)
+        else:
+            self.decoder = None
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, mask, pos_embed, query_embed=None):
+        # flatten NxCxHxW to HWxNxC
+        # src = src.flatten(2).permute(2, 0, 1)
+        # pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
+        # mask = mask.flatten(1)
+        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
+
+        if self.decoder is not None:
+            # query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
+            tgt = torch.zeros_like(query_embed)
+            hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
+                              pos=pos_embed, query_pos=query_embed)
+
+            return hs, memory
+        else:
+            return mask, memory
+
+
+class TransformerEncOnly(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 dim_feedforward=2048, dropout=0.1, activation="relu", normalize_before=False):
+        super().__init__()
+
+        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, mask, pos_embed):
+        # flatten NxCxHxW to HWxNxC
+        src = src.permute(1, 0, 2)
+        pos_embed = pos_embed.permute(1, 0, 2)
+
+        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
+
+        return memory.permute(1, 0,2)
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        output = src
+
+        for layer in self.layers:
+            output = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+
+        intermediate = []
+
+        for layer in self.layers:
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output
+
+
+class TransformerEncoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        """
+            args
+        """
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, src,
+                     src_mask: Optional[Tensor] = None,
+                     src_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward_pre(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+    def forward(self, src,
+                src_mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_transformer(cfg):
+    return Transformer(
+        d_model=cfg.MODEL.VL.HIDDEN_DIM,
+        dropout=cfg.MODEL.VL.DROPOUT,
+        nhead=cfg.MODEL.VL.NHEAD,
+        dim_feedforward=cfg.MODEL.VL.DIM_FEEDFORWARD,
+        num_encoder_layers=cfg.MODEL.VL.ENC_LAYERS,
+        num_decoder_layers=cfg.MODEL.VL.DEC_LAYERS,
+        normalize_before=cfg.MODEL.VL.NORM_BEFORE,
+        return_intermediate_dec=cfg.MODEL.VL.RETURN_INTERMEDIATE,
+        activation=cfg.MODEL.VL.ACTIVATION
+    )
+
+def build_temporal_encoder(cfg):
+    return TransformerEncOnly(d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+                              nhead=8,
+                              num_encoder_layers=2,
+                              dim_feedforward=512, dropout=cfg.MODEL.DECODER.DROPOUT,
+                              activation="GELU", normalize_before=False)
+
+def build_decoder_by_transformer_encoder(cfg):
+    return TransformerEncOnly(d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+                              nhead=8,
+                              num_encoder_layers=6,
+                              dim_feedforward=512, dropout=cfg.MODEL.DECODER.DROPOUT,
+                              activation="GELU", normalize_before=False)
+
+def build_visual_prompt_forward(cfg):
+    return TransformerEncOnly(d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+                              nhead=cfg.MODEL.DECODER.NUM_HEADS,
+                              num_encoder_layers=2,
+                              dim_feedforward=512, dropout=cfg.MODEL.DECODER.DROPOUT,
+                              activation="GELU", normalize_before=False)
+
+def build_text_prompt_forward(cfg):
+    return TransformerEncOnly(d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+                              nhead=cfg.MODEL.DECODER.NUM_HEADS,
+                              num_encoder_layers=2,
+                              dim_feedforward=512, dropout=cfg.MODEL.DECODER.DROPOUT,
+                              activation="GELU", normalize_before=False)
+
+def build_search_delate_template_forward(cfg):
+    return TransformerEncOnly(d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
+                              nhead=cfg.MODEL.DECODER.NUM_HEADS,
+                              num_encoder_layers=2,
+                              dim_feedforward=512, dropout=cfg.MODEL.DECODER.DROPOUT,
+                              activation="GELU", normalize_before=False)
+
+def build_decoder(cfg):
+    decoder_layer = TransformerDecoderLayer(d_model=cfg.MODEL.VL.HIDDEN_DIM, nhead=cfg.MODEL.VL.NHEAD,
+                                            dim_feedforward=cfg.MODEL.VL.DIM_FEEDFORWARD,
+                                            dropout=cfg.MODEL.VL.DROPOUT, activation=cfg.MODEL.VL.ACTIVATION,
+                                            normalize_before=cfg.MODEL.VL.NORM_BEFORE)
+    decoder_norm = nn.LayerNorm(cfg.MODEL.VL.HIDDEN_DIM)
+
+    return TransformerDecoder(decoder_layer, cfg.MODEL.VL.DEC_LAYERS, decoder_norm,
+                              return_intermediate=cfg.MODEL.VL.RETURN_INTERMEDIATE)
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "RELU":
+        return F.relu
+    if activation == "GELU":
+        return F.gelu
+    if activation == "GLU":
+        return F.glu
+    elif activation == "QUICK_GELU":
+        return QuickGELU()
+    raise RuntimeError(F"activation should be RELU/GELU/GLU/QUICK_GELU, not {activation}.")