add visualisation code for attention maps

README.md

@@ -25,6 +25,9 @@ If you use this code for a paper please cite:
 }
 ```
 
+## Updates
+
+- 10/03/2022: Add visualisation code for attention maps in Figure 3.
 
 
 ## Usage
@@ -112,13 +115,50 @@ We provide a script for visualising the learned offsets by the proposed deformab
 conda activate lit
 cd classification/code_for_lit_ti
 
-# visualize
+# visualise
 python visualize_offset.py --model lit_ti --resume [path/to/lit_ti.pth] --vis_image visualization/demo.JPEG
 ```
 
 The plots will be automatically saved under `visualization/`, with a folder named by the name of the example image.
 
 
+## Attention Map Visualisation
+
+We provide our method for visualising the attention maps in Figure 3. To save your time, we also provide the pretrained model for PVT with standard MSA in all stages.
+
+| Name       | Params (M) | FLOPs (G) | Top-1 Acc. (%) | Model  | Log  |
+| ---------- | ---------- | --------- | -------------- | ------ | ---- |
+| PVT w/ MSA | 20         | 8.4       | 80.9           | github | log  |
+
+```bash
+conda activate lit
+cd classification/code_for_lit_ti
+
+# visualise
+# by default, we save the results under 'classification/code_for_lit_ti/attn_results'
+python generate_attention_maps.py --data-path [/path/to/imagenet] --resume [/path/to/pvt_full_msa.pth]
+```
+
+The resulting folder contains the following items,
+
+```
+.
+├── attention_map
+│   ├── stage-0
+│   │   ├── block0
+│   │   │   └── pixel-1260-block-0-head-0.png
+│   │   ├── block1
+│   │   │   └── pixel-1260-block-1-head-0.png
+│   │   └── block2
+│   │       └── pixel-1260-block-2-head-0.png
+│   ├── stage-1
+│   ├── stage-2
+│   └── stage-3
+└── full_msa_eval_maps.npy
+```
+
+where `full_msa_eval_maps.npy` contains the saved attention maps in each block and each stage.  The folder `attention_map` contains the visualisation results.
+
 
 ## License
 

classification/code_for_lit_ti/generate_attention_maps.py

@@ -0,0 +1,160 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import utils
+from timm.models import create_model
+import pvt_full_msa
+import os
+import torch
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import seaborn as sns
+sns.set_theme()
+cmap = sns.light_palette("#A20000", as_cmap=True)
+from params import args
+import torch.backends.cudnn as cudnn
+from datasets import build_dataset
+
+@torch.no_grad()
+def get_attention_data(data_loader, model, device, base_path):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    attention_store = []
+    samples = 0
+    for images, target in metric_logger.log_every(data_loader, 10, header):
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+        samples += images.size()[0]
+        # compute output
+        with torch.cuda.amp.autocast():
+            output, attention_maps = model(images)
+            loss = criterion(output, target)
+            if len(attention_store) == 0:
+                for i, stage_maps in enumerate(attention_maps):
+                    stage_attns = []
+                    for j, block_maps in enumerate(stage_maps):
+                        # Simply use a summation to aggregate the attention probabilities from all batches,
+                        # you can also try to use average or some other scaling methods
+                        stage_attns.append(block_maps.sum(dim=0))
+                    attention_store.append(stage_attns)
+            else:
+                for i, stage_maps in enumerate(attention_maps):
+                    for j, block_maps in enumerate(stage_maps):
+                        attention_store[i][j] += block_maps.sum(dim=0)
+
+        np_attns = []
+        for i, stage_maps in enumerate(attention_store):
+            stage_attns = []
+            for j, block_maps in enumerate(stage_maps):
+                block_maps /= samples
+                stage_attns.append(block_maps.numpy())
+            np_attns.append(stage_attns)
+        np.save(os.path.join(base_path, 'full_msa_eval_maps.npy'), np.array(np_attns))
+        break
+
+def visualize_attentions(base_path):
+
+    save_path = os.path.join(base_path, 'attention_map')
+    attention_maps = np.load(os.path.join(base_path, 'full_msa_eval_maps.npy'), allow_pickle=True)
+
+    linewidths = [1, 1, 2, 2]
+    # Remember that PVT has 4 stages
+    for stage_id, stage_attn_map in enumerate(attention_maps):
+        # each stage has several Transformer blocks
+        for block_id, block_attn_map in enumerate(stage_attn_map):
+
+            block_attn_map = torch.from_numpy(block_attn_map) # size: num_head * seq_len * seq_len
+
+            # PVT has the CLS token at the last stage, here we exclude it for better visualization.
+            if stage_id == 3:
+                test = block_attn_map[:, 1:, :]
+                block_attn_map = test[:, :, 1:]
+
+            H, N, _ = block_attn_map.shape
+            width = int(N ** (1 / 2))
+
+            # iterate each self-attention head
+            for head_id in range(H):
+                head_atth_map = block_attn_map[head_id, ...]
+                map_save_dir = os.path.join(save_path, 'stage-'+str(stage_id), 'block'+str(block_id))
+
+                if not os.path.exists(map_save_dir):
+                    os.makedirs(map_save_dir, exist_ok=True)
+
+                for pixel_id in range(N):
+                    # some random pixel indices, just want to make sure the visualized pixel is near the centre.
+                    if stage_id == 0 and pixel_id != 1260:
+                        continue
+                    if stage_id == 1 and pixel_id != 294:
+                        continue
+                    if stage_id == 2 and pixel_id != 92:
+                        continue
+                    if stage_id == 3 and pixel_id != 17:
+                        continue
+
+                    plt.clf()
+                    f, ax = plt.subplots(1, 1, figsize=(4, 4))
+                    ax.set_aspect('equal')
+
+                    print(stage_id, block_id, head_id, pixel_id)
+
+                    pixel_attn_map = head_atth_map[pixel_id, ...].reshape(int(N ** (1 / 2)), int(N ** (1 / 2))).numpy()
+
+                    x = int(pixel_id % width)
+                    y = int(pixel_id / width)
+
+                    # visualize the attention map with seaborn heatmap
+                    ax = sns.heatmap(pixel_attn_map, cmap="OrRd", cbar=False, xticklabels=False, yticklabels=False, ax=ax)
+                    patch = patches.Rectangle((x, y), 1, 1, linewidth=linewidths[stage_id], edgecolor='lime', facecolor='None')
+                    ax.add_patch(patch)
+                    image_name = 'pixel-{}-block-{}-head-{}.png'.format(pixel_id, block_id, head_id)
+                    plt.savefig(os.path.join(map_save_dir, image_name), transparent=True)
+
+
+
+if __name__ == '__main__':
+    # You may change the path for saving the results.
+    save_path = 'attn_results'
+    if not os.path.exists(save_path):
+        os.makedirs(save_path, exist_ok=True)
+
+    dataset_val, _ = build_dataset(is_train=False, args=args)
+    device = torch.device(args.device)
+
+    # fix the seed for reproducibility
+    seed = args.seed + utils.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+
+    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val, sampler=sampler_val,
+        batch_size=100,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=False
+    )
+
+    model = create_model(
+        'pvt_small_full_msa',
+        pretrained=False,
+        num_classes=1000,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=None,
+    )
+    model.to(device)
+    checkpoint = torch.load(args.resume)
+    model.load_state_dict(checkpoint)
+    get_attention_data(data_loader_val, model, device, save_path)
+    visualize_attentions(save_path)
+
+
+