add fsa

Author: liningping

README.md

@@ -1,2 +1,54 @@
-# InhibitoryAttention
-类脑组抑制性神经元与Attention
+# Win-Take-All Self-Attention
+
+This folder contains the implementation of Win-Take-All Self-Attention based on DeiT and PVT models for image classification.
+
+## Dependencies
+
+- Python 3.9
+- PyTorch == 1.11.0
+- torchvision == 0.12.0
+- numpy
+- timm == 0.4.12
+- einops
+- yacs
+
+```
+pipx install uv
+uv sync
+```
+
+## Data preparation
+
+The ImageNet dataset should be prepared as follows:
+
+```
+$ tree data
+imagenet
+├── train
+│   ├── class1
+│   │   ├── img1.jpeg
+│   │   ├── img2.jpeg
+│   │   └── ...
+│   ├── class2
+│   │   ├── img3.jpeg
+│   │   └── ...
+│   └── ...
+└── val
+    ├── class1
+    │   ├── img4.jpeg
+    │   ├── img5.jpeg
+    │   └── ...
+    ├── class2
+    │   ├── img6.jpeg
+    │   └── ...
+    └── ...
+```
+
+## Train Models from Scratch
+
+- To train `FSA-DeiT/FSA-PVT` on ImageNet from scratch, run:
+
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 main.py --cfg <path-to-config-file> --data-path <imagenet-path> --output <output-path>
+```
+

cfgs/fsa_deit_b.yaml

@@ -0,0 +1,4 @@
+MODEL:
+  TYPE: fsa_deit_base
+  NAME: fsa_deit_base
+  DROP_PATH_RATE: 0.3

cfgs/fsa_deit_s.yaml

@@ -0,0 +1,4 @@
+MODEL:
+  TYPE: fsa_deit_small
+  NAME: fsa_deit_small
+  DROP_PATH_RATE: 0.1

cfgs/fsa_deit_t.yaml

@@ -0,0 +1,6 @@
+DATA:
+  BATCH_SIZE: 512
+MODEL:
+  TYPE: fsa_deit_tiny
+  NAME: fsa_deit_tiny
+  DROP_PATH_RATE: 0.0

cfgs/fsa_pvt_b.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_SIZE: 224
+  BATCH_SIZE: 64
+
+TRAIN:
+  WEIGHT_DECAY: 0.05
+  EPOCHS: 300
+  WARMUP_EPOCHS: 5
+  COOLDOWN_EPOCHS: 10
+  BASE_LR: 5e-4
+  WARMUP_LR: 1e-6
+  MIN_LR: 1e-5
+  CLIP_GRAD: 1.0
+
+MODEL:
+  TYPE: fsa_pvt_large
+  NAME: fsa_pvt_large
+  DROP_PATH_RATE: 0.3
+  FSA:
+    ATTN_TYPE: FFFF

cfgs/fsa_pvt_m.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_SIZE: 224
+  BATCH_SIZE: 128
+
+TRAIN:
+  WEIGHT_DECAY: 0.05
+  EPOCHS: 300
+  WARMUP_EPOCHS: 5
+  COOLDOWN_EPOCHS: 10
+  BASE_LR: 5e-4
+  WARMUP_LR: 1e-6
+  MIN_LR: 1e-5
+  CLIP_GRAD: 1.0
+
+MODEL:
+  TYPE: fsa_pvt_medium
+  NAME: fsa_pvt_medium
+  DROP_PATH_RATE: 0.3
+  FSA:
+    ATTN_TYPE: FFFF

cfgs/fsa_pvt_s.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_SIZE: 224
+  BATCH_SIZE: 128
+
+TRAIN:
+  WEIGHT_DECAY: 0.05
+  EPOCHS: 300
+  WARMUP_EPOCHS: 5
+  COOLDOWN_EPOCHS: 10
+  BASE_LR: 5e-4
+  WARMUP_LR: 1e-6
+  MIN_LR: 1e-5
+  CLIP_GRAD: None
+
+MODEL:
+  TYPE: fsa_pvt_small
+  NAME: fsa_pvt_small
+  DROP_PATH_RATE: 0.1
+  FSA:
+    ATTN_TYPE: FFFF

cfgs/fsa_pvt_t.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_SIZE: 224
+  BATCH_SIZE: 128
+
+TRAIN:
+  WEIGHT_DECAY: 0.05
+  EPOCHS: 300
+  WARMUP_EPOCHS: 5
+  COOLDOWN_EPOCHS: 10
+  BASE_LR: 5e-4
+  WARMUP_LR: 1e-6
+  MIN_LR: 1e-5
+  CLIP_GRAD: None
+
+MODEL:
+  TYPE: fsa_pvt_tiny
+  NAME: fsa_pvt_tiny
+  DROP_PATH_RATE: 0.1
+  FSA:
+    ATTN_TYPE: FFFF