chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+131
@@ -0,0 +1,131 @@
|
||||
# [DiT: Self-Supervised Pre-Training for Document Image Transformer](https://arxiv.org/abs/2203.02378)
|
||||
|
||||
DiT (Document Image Transformer) is a self-supervised pre-trained Document Image Transformer model using large-scale unlabeled text images for Document AI tasks, which is essential since no supervised counterparts ever exist due to the lack of human labeled document images.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/45008728/157173825-0949218a-61f5-4acb-949b-bbc499ab49f2.png" width="500" /><img src="https://user-images.githubusercontent.com/45008728/157173843-796dc878-2607-48d7-85cb-f54a2c007687.png" width="500"/> Model outputs with PubLayNet (left) and ICDAR 2019 cTDaR (right)
|
||||
</div>
|
||||
|
||||
## What's New
|
||||
- Demos on HuggingFace: [Document Layout Analysis](https://huggingface.co/spaces/nielsr/dit-document-layout-analysis), [Document Image Classification](https://huggingface.co/spaces/microsoft/document-image-transformer)
|
||||
- March 2022: release pre-trained [checkpoints](https://huggingface.co/HYPJUDY/dit/tree/main/dit-pts) and fine-tuning [checkpoints](https://huggingface.co/HYPJUDY/dit/tree/main/dit-fts) & codes (DiT-base and DiT-large)
|
||||
- March 2022: release preprint in [arXiv](https://arxiv.org/abs/2203.02378)
|
||||
|
||||
## [Pretrained models](https://huggingface.co/HYPJUDY/dit/tree/main/dit-pts)
|
||||
|
||||
We provide two DiT weights pretrained on [IIT-CDIP Test Collection 1.0](https://dl.acm.org/doi/10.1145/1148170.1148307). The models were pretrained with 224x224 resolution.
|
||||
|
||||
- `DiT-base`: #layer=12; hidden=768; FFN factor=4x; #head=12; patch=16x16 (#parameters: 86M)
|
||||
- `DiT-large`: #layer=24; hidden=1024; FFN factor=4x; #head=16; patch=16x16 (#parameters: 304M)
|
||||
|
||||
Download checkpoints that are **self-supervised pretrained** on IIT-CDIP Test Collection 1.0:
|
||||
- DiT-base: [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth)
|
||||
- DiT-large: [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth)
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
First, clone the repo and install required packages:
|
||||
```
|
||||
git clone https://github.com/microsoft/unilm.git
|
||||
cd unilm/dit
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
The required packages including: [Pytorch](https://pytorch.org/) version 1.9.0, [torchvision](https://pytorch.org/vision/stable/index.html) version 0.10.0 and [Timm](https://github.com/rwightman/pytorch-image-models) version 0.5.4, etc.
|
||||
|
||||
For mixed-precision training, please install [apex](https://github.com/NVIDIA/apex)
|
||||
```
|
||||
git clone https://github.com/NVIDIA/apex
|
||||
cd apex
|
||||
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
|
||||
```
|
||||
For object detection, please additionally install detectron2 library and shapely. Refer to the [Detectron2's INSTALL.md](https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md).
|
||||
|
||||
```bash
|
||||
# Install `detectron2`
|
||||
python -m pip install detectron2 -f \
|
||||
https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.9/index.html
|
||||
|
||||
# Install `shapely`
|
||||
pip install shapely
|
||||
```
|
||||
|
||||
## Fine-tuning on RVL-CDIP (Document Image Classification)
|
||||
|
||||
We summarize the validation results as follows. We also provide the fine-tuned weights. The detailed instructions to reproduce the results can be found at [`classification/README.md`](classification/README.md).
|
||||
|
||||
| name | initialized checkpoint | resolution | accuracy | weight |
|
||||
|------------|:----------------------------------------|:----------:|:-------:|-----|
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | 224x224 | 92.11 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/rvlcdip_dit-b.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | 224x224 | 92.69 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/rvlcdip_dit-l.pth) |
|
||||
|
||||
|
||||
## Fine-tuning on PubLayNet (Document Layout Analysis)
|
||||
|
||||
We summarize the validation results as follows. We also provide the fine-tuned weights. The detailed instructions to reproduce the results can be found at [`object_detection/README.md`](object_detection/README.md).
|
||||
|
||||
| name | initialized checkpoint | detection algorithm | mAP| weight |
|
||||
|------------|:----------------------------------------|:----------:|-------------------|-----|
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Mask R-CNN | 0.935 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/publaynet_dit-b_mrcnn.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Mask R-CNN | 0.941 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/publaynet_dit-l_mrcnn.pth) |
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Cascade R-CNN | 0.945 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/publaynet_dit-b_cascade.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Cascade R-CNN | 0.949 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/publaynet_dit-l_cascade.pth) |
|
||||
|
||||
## Fine-tuning on ICDAR 2019 cTDaR (Table Detection)
|
||||
|
||||
We summarize the validation results as follows. We also provide the fine-tuned weights. The detailed instructions to reproduce the results can be found at [`object_detection/README.md`](object_detection/README.md).
|
||||
|
||||
**Modern**
|
||||
|
||||
| name | initialized checkpoint | detection algorithm | Weighted Average F1 | weight |
|
||||
|------------|:----------------------------------------|:----------:|-------------------|-----|
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Mask R-CNN | 94.74 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19modern_dit-b_mrcnn.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Mask R-CNN | 95.50 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19modern_dit-l_mrcnn.pth) |
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Cascade R-CNN | 95.85 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19modern_dit-b_cascade.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Cascade R-CNN | 96.29 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19modern_dit-l_cascade.pth) |
|
||||
|
||||
**Archival**
|
||||
|
||||
| name | initialized checkpoint | detection algorithm | Weighted Average F1 | weight |
|
||||
|------------|:----------------------------------------|:----------:|-------------------|-----|
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Mask R-CNN | 96.24 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19archival_dit-b_mrcnn.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Mask R-CNN | 96.46 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19archival_dit-l_mrcnn.pth) |
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Cascade R-CNN | 96.63 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19archival_dit-b_cascade.pth) |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Cascade R-CNN | 97.00 | [link](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-fts/icdar19archival_dit-l_cascade.pth) |
|
||||
|
||||
**Combined (Combine the inference results of Modern and Archival)**
|
||||
|
||||
| name | initialized checkpoint | detection algorithm | Weighted Average F1 | weight |
|
||||
|------------|:----------------------------------------|:----------:|-------------------|-----|
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Mask R-CNN | 95.30 | - |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Mask R-CNN | 95.85 | - |
|
||||
| DiT-base | [dit_base_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-base-224-p16-500k-62d53a.pth) | Cascade R-CNN | 96.14 | - |
|
||||
| DiT-large | [dit_large_patch16_224](https://huggingface.co/HYPJUDY/dit/resolve/main/dit-pts/dit-large-224-p16-500k-d7a2fb.pth) | Cascade R-CNN | 96.55 | - |
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider citing our work:
|
||||
```
|
||||
@misc{li2022dit,
|
||||
title={DiT: Self-supervised Pre-training for Document Image Transformer},
|
||||
author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
|
||||
year={2022},
|
||||
eprint={2203.02378},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
This repository is built using the [timm](https://github.com/rwightman/pytorch-image-models) library, the [detectron2](https://github.com/facebookresearch/detectron2) library, the [DeiT](https://github.com/facebookresearch/deit) repository, the [Dino](https://github.com/facebookresearch/dino) repository, the [BEiT](https://github.com/microsoft/unilm/tree/master/beit) repository and the [MPViT](https://github.com/youngwanLEE/MPViT) repository.
|
||||
|
||||
|
||||
### Contact Information
|
||||
|
||||
For help or issues using DiT models, please submit a GitHub issue.
|
||||
|
||||
For other communications related to DiT, please contact Lei Cui (`lecu@microsoft.com`), Furu Wei (`fuwei@microsoft.com`).
|
||||
@@ -0,0 +1,79 @@
|
||||
# DiT for Image Classification
|
||||
|
||||
This folder contains the image classification running instructions on DiT for RVL-CDIP.
|
||||
|
||||
## Usage
|
||||
### Data Preparation
|
||||
|
||||
**RVL-CDIP**
|
||||
|
||||
Download the "rvl-cdip.tar.gz" from this [link](https://www.cs.ryerson.ca/~aharley/rvl-cdip/) (~37GB). Then extract it to `PATH-to-rvlcdip`.
|
||||
|
||||
### Evaluation
|
||||
Following commands provide example to evaluate the fine-tuned checkpoints.
|
||||
```bash
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --master_port=47770 run_class_finetuning.py \
|
||||
--model beit_base_patch16_224 #beit_base_patch16_224 / beit_large_patch16_224
|
||||
--data_path "/path/to/rvlcdip"
|
||||
--eval_data_path "/path/to/rvlcdip"
|
||||
--enable_deepspeed
|
||||
--nb_classes 16
|
||||
--eval
|
||||
--data_set rvlcdip
|
||||
--finetune /path/to/model.pth
|
||||
--output_dir output_dir
|
||||
--log_dir output_dir/tf
|
||||
--batch_size 256
|
||||
--abs_pos_emb
|
||||
--disable_rel_pos_bias
|
||||
```
|
||||
|
||||
### Training
|
||||
Fine-tune DiT on RVL-CDIP:
|
||||
```bash
|
||||
exp_name=dit-base-exp
|
||||
|
||||
mkdir -p output/${exp_name}
|
||||
python -m torch.distributed.launch --nproc_per_node=8 run_class_finetuning.py
|
||||
--model beit_base_patch16_224 #beit_base_patch16_224 / beit_large_patch16_224
|
||||
--data_path "/path/to/rvlcdip"
|
||||
--eval_data_path "/path/to/rvlcdip"
|
||||
--nb_classes 16
|
||||
--data_set rvlcdip
|
||||
--finetune /path/to/model.pth
|
||||
--output_dir output/${exp_name}/
|
||||
--log_dir output/${exp_name}/tf
|
||||
--batch_size 64
|
||||
--lr 5e-4
|
||||
--update_freq 2
|
||||
--eval_freq 10
|
||||
--save_ckpt_freq 10
|
||||
--warmup_epochs 20
|
||||
--epochs 180
|
||||
--layer_scale_init_value 1e-5
|
||||
--layer_decay 0.75
|
||||
--drop_path 0.2
|
||||
--weight_decay 0.05
|
||||
--clip_grad 1.0
|
||||
--abs_pos_emb
|
||||
--disable_rel_pos_bias
|
||||
```
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider citing our work:
|
||||
```
|
||||
@misc{li2022dit,
|
||||
title={DiT: Self-supervised Pre-training for Document Image Transformer},
|
||||
author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
|
||||
year={2022},
|
||||
eprint={2203.02378},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Acknowledgment
|
||||
This part is built using the [timm](https://github.com/rwightman/pytorch-image-models) library, the [Beit](https://github.com/microsoft/unilm/tree/master/beit) repository, the [DeiT](https://github.com/facebookresearch/deit) repository and the [Dino](https://github.com/facebookresearch/dino) repository.
|
||||
@@ -0,0 +1,298 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Modified on torchvision code bases
|
||||
# https://github.com/pytorch/vision
|
||||
# --------------------------------------------------------'
|
||||
import os
|
||||
import os.path
|
||||
import random
|
||||
from typing import Any, Callable, cast, Dict, List, Optional, Tuple
|
||||
|
||||
from PIL import Image
|
||||
from torchvision.datasets.vision import VisionDataset
|
||||
|
||||
|
||||
def has_file_allowed_extension(filename: str, extensions: Tuple[str, ...]) -> bool:
|
||||
"""Checks if a file is an allowed extension.
|
||||
|
||||
Args:
|
||||
filename (string): path to a file
|
||||
extensions (tuple of strings): extensions to consider (lowercase)
|
||||
|
||||
Returns:
|
||||
bool: True if the filename ends with one of given extensions
|
||||
"""
|
||||
return filename.lower().endswith(extensions)
|
||||
|
||||
|
||||
def is_image_file(filename: str) -> bool:
|
||||
"""Checks if a file is an allowed image extension.
|
||||
|
||||
Args:
|
||||
filename (string): path to a file
|
||||
|
||||
Returns:
|
||||
bool: True if the filename ends with a known image extension
|
||||
"""
|
||||
return has_file_allowed_extension(filename, IMG_EXTENSIONS)
|
||||
|
||||
|
||||
def make_dataset(
|
||||
directory: str,
|
||||
class_to_idx: Dict[str, int],
|
||||
extensions: Optional[Tuple[str, ...]] = None,
|
||||
is_valid_file: Optional[Callable[[str], bool]] = None,
|
||||
) -> List[Tuple[str, int]]:
|
||||
instances = []
|
||||
directory = os.path.expanduser(directory)
|
||||
both_none = extensions is None and is_valid_file is None
|
||||
both_something = extensions is not None and is_valid_file is not None
|
||||
if both_none or both_something:
|
||||
raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")
|
||||
if extensions is not None:
|
||||
def is_valid_file(x: str) -> bool:
|
||||
return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))
|
||||
is_valid_file = cast(Callable[[str], bool], is_valid_file)
|
||||
for target_class in sorted(class_to_idx.keys()):
|
||||
class_index = class_to_idx[target_class]
|
||||
target_dir = os.path.join(directory, target_class)
|
||||
if not os.path.isdir(target_dir):
|
||||
continue
|
||||
for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
|
||||
for fname in sorted(fnames):
|
||||
path = os.path.join(root, fname)
|
||||
if is_valid_file(path):
|
||||
item = path, class_index
|
||||
instances.append(item)
|
||||
return instances
|
||||
|
||||
|
||||
class DatasetFolder(VisionDataset):
|
||||
"""A generic data loader where the samples are arranged in this way: ::
|
||||
|
||||
root/class_x/xxx.ext
|
||||
root/class_x/xxy.ext
|
||||
root/class_x/xxz.ext
|
||||
|
||||
root/class_y/123.ext
|
||||
root/class_y/nsdf3.ext
|
||||
root/class_y/asd932_.ext
|
||||
|
||||
Args:
|
||||
root (string): Root directory path.
|
||||
loader (callable): A function to load a sample given its path.
|
||||
extensions (tuple[string]): A list of allowed extensions.
|
||||
both extensions and is_valid_file should not be passed.
|
||||
transform (callable, optional): A function/transform that takes in
|
||||
a sample and returns a transformed version.
|
||||
E.g, ``transforms.RandomCrop`` for images.
|
||||
target_transform (callable, optional): A function/transform that takes
|
||||
in the target and transforms it.
|
||||
is_valid_file (callable, optional): A function that takes path of a file
|
||||
and check if the file is a valid file (used to check of corrupt files)
|
||||
both extensions and is_valid_file should not be passed.
|
||||
|
||||
Attributes:
|
||||
classes (list): List of the class names sorted alphabetically.
|
||||
class_to_idx (dict): Dict with items (class_name, class_index).
|
||||
samples (list): List of (sample path, class_index) tuples
|
||||
targets (list): The class_index value for each image in the dataset
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str,
|
||||
loader: Callable[[str], Any],
|
||||
extensions: Optional[Tuple[str, ...]] = None,
|
||||
transform: Optional[Callable] = None,
|
||||
target_transform: Optional[Callable] = None,
|
||||
is_valid_file: Optional[Callable[[str], bool]] = None,
|
||||
) -> None:
|
||||
super(DatasetFolder, self).__init__(root, transform=transform,
|
||||
target_transform=target_transform)
|
||||
classes, class_to_idx = self._find_classes(self.root)
|
||||
samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file)
|
||||
if len(samples) == 0:
|
||||
msg = "Found 0 files in subfolders of: {}\n".format(self.root)
|
||||
if extensions is not None:
|
||||
msg += "Supported extensions are: {}".format(",".join(extensions))
|
||||
raise RuntimeError(msg)
|
||||
|
||||
self.loader = loader
|
||||
self.extensions = extensions
|
||||
|
||||
self.classes = classes
|
||||
self.class_to_idx = class_to_idx
|
||||
self.samples = samples
|
||||
self.targets = [s[1] for s in samples]
|
||||
|
||||
def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
|
||||
"""
|
||||
Finds the class folders in a dataset.
|
||||
|
||||
Args:
|
||||
dir (string): Root directory path.
|
||||
|
||||
Returns:
|
||||
tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
|
||||
|
||||
Ensures:
|
||||
No class is a subdirectory of another.
|
||||
"""
|
||||
classes = [d.name for d in os.scandir(dir) if d.is_dir()]
|
||||
classes.sort()
|
||||
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
|
||||
return classes, class_to_idx
|
||||
|
||||
def __getitem__(self, index: int) -> Tuple[Any, Any]:
|
||||
"""
|
||||
Args:
|
||||
index (int): Index
|
||||
|
||||
Returns:
|
||||
tuple: (sample, target) where target is class_index of the target class.
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
path, target = self.samples[index]
|
||||
sample = self.loader(path)
|
||||
break
|
||||
except Exception as e:
|
||||
print(e)
|
||||
index = random.randint(0, len(self.samples) - 1)
|
||||
|
||||
if self.transform is not None:
|
||||
sample = self.transform(sample)
|
||||
if self.target_transform is not None:
|
||||
target = self.target_transform(target)
|
||||
|
||||
return sample, target
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
|
||||
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
|
||||
|
||||
|
||||
def pil_loader(path: str) -> Image.Image:
|
||||
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
|
||||
with open(path, 'rb') as f:
|
||||
img = Image.open(f)
|
||||
return img.convert('RGB')
|
||||
|
||||
|
||||
# TODO: specify the return type
|
||||
def accimage_loader(path: str) -> Any:
|
||||
import accimage
|
||||
try:
|
||||
return accimage.Image(path)
|
||||
except IOError:
|
||||
# Potentially a decoding problem, fall back to PIL.Image
|
||||
return pil_loader(path)
|
||||
|
||||
|
||||
def default_loader(path: str) -> Any:
|
||||
from torchvision import get_image_backend
|
||||
if get_image_backend() == 'accimage':
|
||||
return accimage_loader(path)
|
||||
else:
|
||||
return pil_loader(path)
|
||||
|
||||
|
||||
class RvlcdipDatasetFolder(VisionDataset):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str,
|
||||
loader: Callable[[str], Any],
|
||||
extensions: Optional[Tuple[str, ...]] = None,
|
||||
transform: Optional[Callable] = None,
|
||||
target_transform: Optional[Callable] = None,
|
||||
split: str = None,
|
||||
dataset_size: Optional[int] = None
|
||||
) -> None:
|
||||
super().__init__(root, transform=transform, target_transform=target_transform)
|
||||
self.dataset_size = int(dataset_size) if dataset_size is not None else 42948004
|
||||
classes = ["letter",
|
||||
"form",
|
||||
"email",
|
||||
"handwritten",
|
||||
"advertisement",
|
||||
"scientific report",
|
||||
"scientific publication",
|
||||
"specification",
|
||||
"file folder",
|
||||
"news article",
|
||||
"budget",
|
||||
"invoice",
|
||||
"presentation",
|
||||
"questionnaire",
|
||||
"resume",
|
||||
"memo"]
|
||||
class_to_idx = {c: i for i, c in enumerate(classes)}
|
||||
with open(os.path.join(self.root, "labels", split + ".txt"), "r") as f:
|
||||
labels = f.read().splitlines()
|
||||
samples = [(line.split()[0], int(line.split()[1])) for line in labels]
|
||||
try:
|
||||
assert len(samples) > 0 and os.path.exists(os.path.join(self.root, "images", samples[0][0]))
|
||||
except:
|
||||
msg = "Found 0 files in subfolders of: {}\n".format(self.root)
|
||||
msg += "Expected first file: {}".format(os.path.join(self.root, "images", samples[0][0]))
|
||||
raise RuntimeError(msg)
|
||||
|
||||
self.loader = loader
|
||||
self.extensions = extensions
|
||||
|
||||
self.classes = classes
|
||||
self.class_to_idx = class_to_idx
|
||||
self.samples = samples
|
||||
self.targets = [s[1] for s in samples]
|
||||
|
||||
def __getitem__(self, index: int) -> Tuple[Any, Any]:
|
||||
"""
|
||||
Args:
|
||||
index (int): Index
|
||||
Returns:
|
||||
tuple: (sample, target) where target is class_index of the target class.
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
path, target = self.samples[index]
|
||||
sample = self.loader(os.path.join(self.root, "images", path))
|
||||
break
|
||||
except Exception as e:
|
||||
print(e)
|
||||
index = random.randint(0, len(self.samples) - 1)
|
||||
|
||||
if self.transform is not None:
|
||||
sample = self.transform(sample)
|
||||
if self.target_transform is not None:
|
||||
target = self.target_transform(target)
|
||||
|
||||
return sample, target
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
|
||||
class RvlcdipImageFolder(RvlcdipDatasetFolder):
|
||||
def __init__(
|
||||
self,
|
||||
root: str,
|
||||
transform: Optional[Callable] = None,
|
||||
target_transform: Optional[Callable] = None,
|
||||
loader: Callable[[str], Any] = default_loader,
|
||||
split: str = None,
|
||||
dataset_size: Optional[int] = None
|
||||
):
|
||||
super().__init__(root, loader, IMG_EXTENSIONS if split is None else None,
|
||||
transform=transform,
|
||||
target_transform=target_transform,
|
||||
split=split,
|
||||
dataset_size=dataset_size)
|
||||
self.imgs = self.samples
|
||||
@@ -0,0 +1,92 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Based on timm, DINO 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 timm.data import create_transform
|
||||
from timm.data.constants import \
|
||||
IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
|
||||
from timm.data.transforms import str_to_interp_mode
|
||||
from torchvision import transforms
|
||||
|
||||
from dataset_folder import RvlcdipImageFolder
|
||||
|
||||
|
||||
def build_dataset(is_train, args):
|
||||
transform = build_transform(is_train, args)
|
||||
|
||||
print("Transform = ")
|
||||
if isinstance(transform, tuple):
|
||||
for trans in transform:
|
||||
print(" - - - - - - - - - - ")
|
||||
for t in trans.transforms:
|
||||
print(t)
|
||||
else:
|
||||
for t in transform.transforms:
|
||||
print(t)
|
||||
print("---------------------------")
|
||||
|
||||
if args.data_set == 'rvlcdip':
|
||||
root = args.data_path if is_train else args.eval_data_path
|
||||
split = "train" if is_train else "test"
|
||||
dataset = RvlcdipImageFolder(root, split=split, transform=transform)
|
||||
nb_classes = args.nb_classes
|
||||
assert len(dataset.class_to_idx) == nb_classes
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
assert nb_classes == args.nb_classes
|
||||
print("Number of the class = %d" % args.nb_classes)
|
||||
|
||||
return dataset, nb_classes
|
||||
|
||||
|
||||
def build_transform(is_train, args):
|
||||
resize_im = args.input_size > 32
|
||||
imagenet_default_mean_and_std = args.imagenet_default_mean_and_std
|
||||
mean = IMAGENET_INCEPTION_MEAN if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_MEAN
|
||||
std = IMAGENET_INCEPTION_STD if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_STD
|
||||
|
||||
if is_train:
|
||||
# this should always dispatch to transforms_imagenet_train
|
||||
transform = create_transform(
|
||||
input_size=args.input_size,
|
||||
is_training=True,
|
||||
color_jitter=args.color_jitter,
|
||||
auto_augment=args.aa,
|
||||
interpolation=args.train_interpolation,
|
||||
re_prob=args.reprob,
|
||||
re_mode=args.remode,
|
||||
re_count=args.recount,
|
||||
mean=mean,
|
||||
std=std,
|
||||
)
|
||||
if not resize_im:
|
||||
# replace RandomResizedCropAndInterpolation with
|
||||
# RandomCrop
|
||||
transform.transforms[0] = transforms.RandomCrop(
|
||||
args.input_size, padding=4)
|
||||
return transform
|
||||
|
||||
t = []
|
||||
if resize_im:
|
||||
if args.crop_pct is None:
|
||||
if args.input_size < 384:
|
||||
args.crop_pct = 224 / 256
|
||||
else:
|
||||
args.crop_pct = 1.0
|
||||
size = int(args.input_size / args.crop_pct)
|
||||
t.append(
|
||||
transforms.Resize(size, interpolation=str_to_interp_mode("bicubic")),
|
||||
# to maintain same ratio w.r.t. 224 images
|
||||
)
|
||||
t.append(transforms.CenterCrop(args.input_size))
|
||||
|
||||
t.append(transforms.ToTensor())
|
||||
t.append(transforms.Normalize(mean, std))
|
||||
return transforms.Compose(t)
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"train_batch_size": 2048,
|
||||
"train_micro_batch_size_per_gpu": 32,
|
||||
"steps_per_print": 1000,
|
||||
"optimizer": {
|
||||
"type": "Adam",
|
||||
"adam_w_mode": true,
|
||||
"params": {
|
||||
"lr": 0.001,
|
||||
"weight_decay": 0.05,
|
||||
"bias_correction": true,
|
||||
"betas": [
|
||||
0.9,
|
||||
0.999
|
||||
]
|
||||
}
|
||||
},
|
||||
"gradient_clipping": 1,
|
||||
"fp16": {
|
||||
"enabled": true,
|
||||
"loss_scale": 0,
|
||||
"initial_scale_power": 7,
|
||||
"loss_scale_window": 128
|
||||
},
|
||||
"zero_optimization": {
|
||||
"stage": 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Based on timm, DINO and DeiT code bases
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# https://github.com/facebookresearch/deit/
|
||||
# https://github.com/facebookresearch/dino
|
||||
# --------------------------------------------------------'
|
||||
import math
|
||||
import sys
|
||||
from typing import Iterable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from timm.data import Mixup
|
||||
from timm.utils import accuracy, ModelEma
|
||||
|
||||
import utils
|
||||
|
||||
|
||||
def train_class_batch(model, samples, target, criterion):
|
||||
outputs = model(samples)
|
||||
if not isinstance(outputs, torch.Tensor):
|
||||
# assume that the model outputs a tuple of [outputs, outputs_kd]
|
||||
outputs, outputs_kd = outputs
|
||||
loss = criterion(outputs, target)
|
||||
return loss, outputs
|
||||
|
||||
|
||||
def get_loss_scale_for_deepspeed(model):
|
||||
optimizer = model.optimizer
|
||||
return optimizer.loss_scale if hasattr(optimizer, "loss_scale") else optimizer.cur_scale
|
||||
|
||||
|
||||
def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
|
||||
data_loader: Iterable, optimizer: torch.optim.Optimizer,
|
||||
device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
|
||||
model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None, log_writer=None,
|
||||
start_steps=None, lr_schedule_values=None, wd_schedule_values=None,
|
||||
num_training_steps_per_epoch=None, update_freq=None):
|
||||
model.train(True)
|
||||
metric_logger = utils.MetricLogger(delimiter=" ")
|
||||
metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
|
||||
metric_logger.add_meter('min_lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
|
||||
header = 'Epoch: [{}]'.format(epoch)
|
||||
print_freq = 10
|
||||
|
||||
if loss_scaler is None:
|
||||
model.zero_grad()
|
||||
model.micro_steps = 0
|
||||
else:
|
||||
optimizer.zero_grad()
|
||||
|
||||
for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
|
||||
step = data_iter_step // update_freq
|
||||
if step >= num_training_steps_per_epoch:
|
||||
continue
|
||||
it = start_steps + step # global training iteration
|
||||
# Update LR & WD for the first acc
|
||||
if lr_schedule_values is not None or wd_schedule_values is not None and data_iter_step % update_freq == 0:
|
||||
for i, param_group in enumerate(optimizer.param_groups):
|
||||
if lr_schedule_values is not None:
|
||||
param_group["lr"] = lr_schedule_values[it] * param_group["lr_scale"]
|
||||
if wd_schedule_values is not None and param_group["weight_decay"] > 0:
|
||||
param_group["weight_decay"] = wd_schedule_values[it]
|
||||
|
||||
samples = samples.to(device, non_blocking=True)
|
||||
targets = targets.to(device, non_blocking=True)
|
||||
|
||||
if mixup_fn is not None:
|
||||
samples, targets = mixup_fn(samples, targets)
|
||||
|
||||
if loss_scaler is None:
|
||||
samples = samples.half()
|
||||
loss, output = train_class_batch(
|
||||
model, samples, targets, criterion)
|
||||
else:
|
||||
with torch.cuda.amp.autocast():
|
||||
loss, output = train_class_batch(
|
||||
model, samples, targets, criterion)
|
||||
|
||||
loss_value = loss.item()
|
||||
|
||||
if not math.isfinite(loss_value):
|
||||
print("Loss is {}, stopping training".format(loss_value))
|
||||
sys.exit(1)
|
||||
|
||||
if loss_scaler is None:
|
||||
loss /= update_freq
|
||||
model.backward(loss)
|
||||
model.step()
|
||||
|
||||
if (data_iter_step + 1) % update_freq == 0:
|
||||
# model.zero_grad()
|
||||
# Deepspeed will call step() & model.zero_grad() automatic
|
||||
if model_ema is not None:
|
||||
model_ema.update(model)
|
||||
grad_norm = None
|
||||
loss_scale_value = get_loss_scale_for_deepspeed(model)
|
||||
else:
|
||||
# this attribute is added by timm on one optimizer (adahessian)
|
||||
is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
|
||||
loss /= update_freq
|
||||
grad_norm = loss_scaler(loss, optimizer, clip_grad=max_norm,
|
||||
parameters=model.parameters(), create_graph=is_second_order,
|
||||
update_grad=(data_iter_step + 1) % update_freq == 0)
|
||||
if (data_iter_step + 1) % update_freq == 0:
|
||||
optimizer.zero_grad()
|
||||
if model_ema is not None:
|
||||
model_ema.update(model)
|
||||
loss_scale_value = loss_scaler.state_dict()["scale"]
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if mixup_fn is None:
|
||||
class_acc = (output.max(-1)[-1] == targets).float().mean()
|
||||
else:
|
||||
class_acc = None
|
||||
metric_logger.update(loss=loss_value)
|
||||
metric_logger.update(class_acc=class_acc)
|
||||
metric_logger.update(loss_scale=loss_scale_value)
|
||||
min_lr = 10.
|
||||
max_lr = 0.
|
||||
for group in optimizer.param_groups:
|
||||
min_lr = min(min_lr, group["lr"])
|
||||
max_lr = max(max_lr, group["lr"])
|
||||
|
||||
metric_logger.update(lr=max_lr)
|
||||
metric_logger.update(min_lr=min_lr)
|
||||
weight_decay_value = None
|
||||
for group in optimizer.param_groups:
|
||||
if group["weight_decay"] > 0:
|
||||
weight_decay_value = group["weight_decay"]
|
||||
metric_logger.update(weight_decay=weight_decay_value)
|
||||
metric_logger.update(grad_norm=grad_norm)
|
||||
|
||||
if log_writer is not None:
|
||||
log_writer.update(loss=loss_value, head="loss")
|
||||
log_writer.update(class_acc=class_acc, head="loss")
|
||||
log_writer.update(loss_scale=loss_scale_value, head="opt")
|
||||
log_writer.update(lr=max_lr, head="opt")
|
||||
log_writer.update(min_lr=min_lr, head="opt")
|
||||
log_writer.update(weight_decay=weight_decay_value, head="opt")
|
||||
log_writer.update(grad_norm=grad_norm, head="opt")
|
||||
|
||||
log_writer.set_step()
|
||||
|
||||
# gather the stats from all processes
|
||||
metric_logger.synchronize_between_processes()
|
||||
print("Averaged stats:", metric_logger)
|
||||
return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(data_loader, model, device):
|
||||
criterion = torch.nn.CrossEntropyLoss()
|
||||
|
||||
metric_logger = utils.MetricLogger(delimiter=" ")
|
||||
header = 'Test:'
|
||||
|
||||
# switch to evaluation mode
|
||||
model.eval()
|
||||
|
||||
for batch in metric_logger.log_every(data_loader, 10, header):
|
||||
images = batch[0]
|
||||
target = batch[-1]
|
||||
images = images.to(device, non_blocking=True)
|
||||
target = target.to(device, non_blocking=True)
|
||||
|
||||
# compute output
|
||||
with torch.cuda.amp.autocast():
|
||||
output = model(images)
|
||||
loss = criterion(output, target)
|
||||
|
||||
acc1, acc5 = accuracy(output, target, topk=(1, 5))
|
||||
|
||||
batch_size = images.shape[0]
|
||||
metric_logger.update(loss=loss.item())
|
||||
metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
|
||||
metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
|
||||
# gather the stats from all processes
|
||||
metric_logger.synchronize_between_processes()
|
||||
print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
|
||||
.format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
|
||||
|
||||
return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
|
||||
@@ -0,0 +1,409 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# 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
|
||||
# --------------------------------------------------------'
|
||||
import math
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from timm.models.layers import drop_path, 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': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
|
||||
**kwargs
|
||||
}
|
||||
|
||||
|
||||
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.):
|
||||
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)
|
||||
# commit this for the orignal BERT implement
|
||||
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., window_size=None, attn_head_dim=None):
|
||||
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.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
|
||||
|
||||
if window_size:
|
||||
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)
|
||||
else:
|
||||
self.window_size = None
|
||||
self.relative_position_bias_table = None
|
||||
self.relative_position_index = None
|
||||
|
||||
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):
|
||||
B, N, C = x.shape
|
||||
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 = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
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)
|
||||
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 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)
|
||||
|
||||
if rel_pos_bias is not None:
|
||||
attn = attn + rel_pos_bias
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
|
||||
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., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
|
||||
window_size=None, attn_head_dim=None):
|
||||
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, window_size=window_size, attn_head_dim=attn_head_dim)
|
||||
# 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)
|
||||
|
||||
if init_values > 0:
|
||||
self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
|
||||
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, rel_pos_bias=None):
|
||||
if self.gamma_1 is None:
|
||||
x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
|
||||
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
||||
else:
|
||||
x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
|
||||
x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
""" Image to Patch Embedding
|
||||
"""
|
||||
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
|
||||
self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.num_patches = num_patches
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x, **kwargs):
|
||||
B, C, H, W = x.shape
|
||||
# FIXME look at relaxing size constraints
|
||||
assert H == self.img_size[0] and W == self.img_size[1], \
|
||||
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
||||
x = self.proj(x).flatten(2).transpose(1, 2)
|
||||
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)
|
||||
|
||||
# trunc_normal_(self.relative_position_bias_table, std=.02)
|
||||
|
||||
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
|
||||
|
||||
|
||||
class VisionTransformer(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
def __init__(self, img_size=224, 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., norm_layer=nn.LayerNorm, init_values=None,
|
||||
use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False,
|
||||
use_mean_pooling=True, init_scale=0.001):
|
||||
super().__init__()
|
||||
self.num_classes = num_classes
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
# self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
if use_abs_pos_emb:
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, 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
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
self.use_rel_pos_bias = use_rel_pos_bias
|
||||
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,
|
||||
init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)
|
||||
for i in range(depth)])
|
||||
self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
|
||||
self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
|
||||
self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
# trunc_normal_(self.mask_token, std=.02)
|
||||
trunc_normal_(self.head.weight, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
self.head.weight.data.mul_(init_scale)
|
||||
self.head.bias.data.mul_(init_scale)
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
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):
|
||||
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, x):
|
||||
x = self.patch_embed(x)
|
||||
batch_size, seq_len, _ = x.size()
|
||||
|
||||
cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
if self.pos_embed is not None:
|
||||
x = x + self.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:
|
||||
x = blk(x, rel_pos_bias=rel_pos_bias)
|
||||
|
||||
x = self.norm(x)
|
||||
if self.fc_norm is not None:
|
||||
t = x[:, 1:, :]
|
||||
return self.fc_norm(t.mean(1))
|
||||
else:
|
||||
return x[:, 0]
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_small_patch16_224(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
patch_size=16, embed_dim=384, depth=12, num_heads=8, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_base_patch16_224(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_base_patch16_384(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_large_patch16_224(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_large_patch16_384(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
img_size=384, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
|
||||
@register_model
|
||||
def beit_large_patch16_512(pretrained=False, **kwargs):
|
||||
model = VisionTransformer(
|
||||
img_size=512, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
@@ -0,0 +1,179 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Based on timm code bases
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# --------------------------------------------------------'
|
||||
import torch
|
||||
from torch import optim as optim
|
||||
|
||||
from timm.optim.adafactor import Adafactor
|
||||
from timm.optim.adahessian import Adahessian
|
||||
from timm.optim.adamp import AdamP
|
||||
from timm.optim.lookahead import Lookahead
|
||||
from timm.optim.nadam import Nadam
|
||||
from timm.optim.nvnovograd import NvNovoGrad
|
||||
from timm.optim.radam import RAdam
|
||||
from timm.optim.rmsprop_tf import RMSpropTF
|
||||
from timm.optim.sgdp import SGDP
|
||||
|
||||
import json
|
||||
|
||||
try:
|
||||
from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
|
||||
has_apex = True
|
||||
except ImportError:
|
||||
has_apex = False
|
||||
|
||||
|
||||
def get_num_layer_for_vit(var_name, num_max_layer):
|
||||
if var_name in ("cls_token", "mask_token", "pos_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("patch_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("rel_pos_bias"):
|
||||
return num_max_layer - 1
|
||||
elif var_name.startswith("blocks"):
|
||||
layer_id = int(var_name.split('.')[1])
|
||||
return layer_id + 1
|
||||
else:
|
||||
return num_max_layer - 1
|
||||
|
||||
|
||||
class LayerDecayValueAssigner(object):
|
||||
def __init__(self, values):
|
||||
self.values = values
|
||||
|
||||
def get_scale(self, layer_id):
|
||||
return self.values[layer_id]
|
||||
|
||||
def get_layer_id(self, var_name):
|
||||
return get_num_layer_for_vit(var_name, len(self.values))
|
||||
|
||||
|
||||
def get_parameter_groups(model, weight_decay=1e-5, skip_list=(), get_num_layer=None, get_layer_scale=None):
|
||||
parameter_group_names = {}
|
||||
parameter_group_vars = {}
|
||||
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue # frozen weights
|
||||
if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list:
|
||||
group_name = "no_decay"
|
||||
this_weight_decay = 0.
|
||||
else:
|
||||
group_name = "decay"
|
||||
this_weight_decay = weight_decay
|
||||
if get_num_layer is not None:
|
||||
layer_id = get_num_layer(name)
|
||||
group_name = "layer_%d_%s" % (layer_id, group_name)
|
||||
else:
|
||||
layer_id = None
|
||||
|
||||
if group_name not in parameter_group_names:
|
||||
if get_layer_scale is not None:
|
||||
scale = get_layer_scale(layer_id)
|
||||
else:
|
||||
scale = 1.
|
||||
|
||||
parameter_group_names[group_name] = {
|
||||
"weight_decay": this_weight_decay,
|
||||
"params": [],
|
||||
"lr_scale": scale
|
||||
}
|
||||
parameter_group_vars[group_name] = {
|
||||
"weight_decay": this_weight_decay,
|
||||
"params": [],
|
||||
"lr_scale": scale
|
||||
}
|
||||
|
||||
parameter_group_vars[group_name]["params"].append(param)
|
||||
parameter_group_names[group_name]["params"].append(name)
|
||||
print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
|
||||
return list(parameter_group_vars.values())
|
||||
|
||||
|
||||
def create_optimizer(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None):
|
||||
opt_lower = args.opt.lower()
|
||||
weight_decay = args.weight_decay
|
||||
if weight_decay and filter_bias_and_bn:
|
||||
skip = {}
|
||||
if skip_list is not None:
|
||||
skip = skip_list
|
||||
elif hasattr(model, 'no_weight_decay'):
|
||||
skip = model.no_weight_decay()
|
||||
parameters = get_parameter_groups(model, weight_decay, skip, get_num_layer, get_layer_scale)
|
||||
weight_decay = 0.
|
||||
else:
|
||||
parameters = model.parameters()
|
||||
|
||||
if 'fused' in opt_lower:
|
||||
assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers'
|
||||
|
||||
opt_args = dict(lr=args.lr, weight_decay=weight_decay)
|
||||
if hasattr(args, 'opt_eps') and args.opt_eps is not None:
|
||||
opt_args['eps'] = args.opt_eps
|
||||
if hasattr(args, 'opt_betas') and args.opt_betas is not None:
|
||||
opt_args['betas'] = args.opt_betas
|
||||
|
||||
opt_split = opt_lower.split('_')
|
||||
opt_lower = opt_split[-1]
|
||||
if opt_lower == 'sgd' or opt_lower == 'nesterov':
|
||||
opt_args.pop('eps', None)
|
||||
optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
|
||||
elif opt_lower == 'momentum':
|
||||
opt_args.pop('eps', None)
|
||||
optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
|
||||
elif opt_lower == 'adam':
|
||||
optimizer = optim.Adam(parameters, **opt_args)
|
||||
elif opt_lower == 'adamw':
|
||||
optimizer = optim.AdamW(parameters, **opt_args)
|
||||
elif opt_lower == 'nadam':
|
||||
optimizer = Nadam(parameters, **opt_args)
|
||||
elif opt_lower == 'radam':
|
||||
optimizer = RAdam(parameters, **opt_args)
|
||||
elif opt_lower == 'adamp':
|
||||
optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
|
||||
elif opt_lower == 'sgdp':
|
||||
optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
|
||||
elif opt_lower == 'adadelta':
|
||||
optimizer = optim.Adadelta(parameters, **opt_args)
|
||||
elif opt_lower == 'adafactor':
|
||||
if not args.lr:
|
||||
opt_args['lr'] = None
|
||||
optimizer = Adafactor(parameters, **opt_args)
|
||||
elif opt_lower == 'adahessian':
|
||||
optimizer = Adahessian(parameters, **opt_args)
|
||||
elif opt_lower == 'rmsprop':
|
||||
optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
|
||||
elif opt_lower == 'rmsproptf':
|
||||
optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
|
||||
elif opt_lower == 'novograd' or opt_lower == 'nvnovograd':
|
||||
optimizer = NvNovoGrad(parameters, **opt_args)
|
||||
elif opt_lower == 'fusedsgd':
|
||||
opt_args.pop('eps', None)
|
||||
optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
|
||||
elif opt_lower == 'fusedmomentum':
|
||||
opt_args.pop('eps', None)
|
||||
optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
|
||||
elif opt_lower == 'fusedadam':
|
||||
optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
|
||||
elif opt_lower == 'fusedadamw':
|
||||
optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
|
||||
elif opt_lower == 'fusedlamb':
|
||||
optimizer = FusedLAMB(parameters, **opt_args)
|
||||
elif opt_lower == 'fusednovograd':
|
||||
opt_args.setdefault('betas', (0.95, 0.98))
|
||||
optimizer = FusedNovoGrad(parameters, **opt_args)
|
||||
else:
|
||||
assert False and "Invalid optimizer"
|
||||
raise ValueError
|
||||
|
||||
if len(opt_split) > 1:
|
||||
if opt_split[0] == 'lookahead':
|
||||
optimizer = Lookahead(optimizer)
|
||||
|
||||
return optimizer
|
||||
@@ -0,0 +1,16 @@
|
||||
# timm==0.4.12
|
||||
git+https://github.com/rwightman/pytorch-image-models.git
|
||||
taming-transformers-rom1504
|
||||
deepspeed==0.4.0
|
||||
Pillow
|
||||
numpy
|
||||
requests
|
||||
einops
|
||||
tensorboard
|
||||
scipy
|
||||
attrs
|
||||
pybase64
|
||||
pyyaml
|
||||
webdataset
|
||||
omegaconf==2.0.0
|
||||
pytorch-lightning==1.6.0
|
||||
@@ -0,0 +1,615 @@
|
||||
# --------------------------------------------------------
|
||||
# DIT: SELF-SUPERVISED PRE-TRAINING FOR DOCUMENT IMAGE TRANSFORMER
|
||||
# Based on Beit
|
||||
# --------------------------------------------------------'
|
||||
import argparse
|
||||
import datetime
|
||||
import numpy as np
|
||||
import time
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import json
|
||||
import os
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from timm.data.mixup import Mixup
|
||||
from timm.models import create_model
|
||||
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
|
||||
from timm.utils import ModelEma
|
||||
from optim_factory import create_optimizer, get_parameter_groups, LayerDecayValueAssigner
|
||||
import webdataset as wds
|
||||
from datasets import build_dataset
|
||||
from engine_for_finetuning import train_one_epoch, evaluate
|
||||
from utils import NativeScalerWithGradNormCount as NativeScaler
|
||||
import utils
|
||||
from scipy import interpolate
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser('BEiT fine-tuning and evaluation script for image classification', add_help=False)
|
||||
parser.add_argument('--batch_size', default=64, type=int)
|
||||
parser.add_argument('--epochs', default=30, type=int)
|
||||
parser.add_argument('--update_freq', default=1, type=int)
|
||||
parser.add_argument('--save_ckpt_freq', default=5, type=int)
|
||||
parser.add_argument('--eval_freq', default=5, type=int)
|
||||
|
||||
# Model parameters
|
||||
parser.add_argument('--model', default='deit_base_patch16_224', type=str, metavar='MODEL',
|
||||
help='Name of model to train')
|
||||
parser.add_argument('--rel_pos_bias', action='store_true')
|
||||
parser.add_argument('--disable_rel_pos_bias', action='store_false', dest='rel_pos_bias')
|
||||
parser.set_defaults(rel_pos_bias=True)
|
||||
parser.add_argument('--abs_pos_emb', action='store_true')
|
||||
parser.add_argument('--qkv_bias', action='store_true')
|
||||
parser.add_argument('--layer_scale_init_value', default=0.1, type=float,
|
||||
help="0.1 for base, 1e-5 for large. set 0 to disable layer scale")
|
||||
|
||||
parser.add_argument('--input_size', default=224, type=int,
|
||||
help='images input size')
|
||||
|
||||
parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
|
||||
help='Dropout rate (default: 0.)')
|
||||
parser.add_argument('--attn_drop_rate', type=float, default=0.0, metavar='PCT',
|
||||
help='Attention dropout rate (default: 0.)')
|
||||
parser.add_argument('--drop_path', type=float, default=0.1, metavar='PCT',
|
||||
help='Drop path rate (default: 0.1)')
|
||||
|
||||
parser.add_argument('--disable_eval_during_finetuning', action='store_true', default=False)
|
||||
|
||||
parser.add_argument('--model_ema', action='store_true', default=False)
|
||||
parser.add_argument('--model_ema_decay', type=float, default=0.9999, help='')
|
||||
parser.add_argument('--model_ema_force_cpu', action='store_true', default=False, help='')
|
||||
|
||||
# Optimizer parameters
|
||||
parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
|
||||
help='Optimizer (default: "adamw"')
|
||||
parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON',
|
||||
help='Optimizer Epsilon (default: 1e-8)')
|
||||
parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA',
|
||||
help='Optimizer Betas (default: None, use opt default)')
|
||||
parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
|
||||
help='Clip gradient norm (default: None, no clipping)')
|
||||
parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
|
||||
help='SGD momentum (default: 0.9)')
|
||||
parser.add_argument('--weight_decay', type=float, default=0.05,
|
||||
help='weight decay (default: 0.05)')
|
||||
parser.add_argument('--weight_decay_end', type=float, default=None, help="""Final value of the
|
||||
weight decay. We use a cosine schedule for WD and using a larger decay by
|
||||
the end of training improves performance for ViTs.""")
|
||||
|
||||
parser.add_argument('--lr', type=float, default=5e-4, metavar='LR',
|
||||
help='learning rate (default: 5e-4)')
|
||||
parser.add_argument('--layer_decay', type=float, default=0.9)
|
||||
|
||||
parser.add_argument('--warmup_lr', type=float, default=1e-6, metavar='LR',
|
||||
help='warmup learning rate (default: 1e-6)')
|
||||
parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
|
||||
help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
|
||||
|
||||
parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
|
||||
help='epochs to warmup LR, if scheduler supports')
|
||||
parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N',
|
||||
help='num of steps to warmup LR, will overload warmup_epochs if set > 0')
|
||||
|
||||
# Augmentation parameters
|
||||
parser.add_argument('--color_jitter', type=float, default=0.4, metavar='PCT',
|
||||
help='Color jitter factor (default: 0.4)')
|
||||
parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
|
||||
help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)'),
|
||||
parser.add_argument('--smoothing', type=float, default=0.1,
|
||||
help='Label smoothing (default: 0.1)')
|
||||
parser.add_argument('--train_interpolation', type=str, default='bicubic',
|
||||
help='Training interpolation (random, bilinear, bicubic default: "bicubic")')
|
||||
|
||||
# Evaluation parameters
|
||||
parser.add_argument('--crop_pct', type=float, default=None)
|
||||
|
||||
# * Random Erase params
|
||||
parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
|
||||
help='Random erase prob (default: 0.25)')
|
||||
parser.add_argument('--remode', type=str, default='pixel',
|
||||
help='Random erase mode (default: "pixel")')
|
||||
parser.add_argument('--recount', type=int, default=1,
|
||||
help='Random erase count (default: 1)')
|
||||
parser.add_argument('--resplit', action='store_true', default=False,
|
||||
help='Do not random erase first (clean) augmentation split')
|
||||
|
||||
# * Mixup params
|
||||
parser.add_argument('--mixup', type=float, default=0,
|
||||
help='mixup alpha, mixup enabled if > 0.')
|
||||
parser.add_argument('--cutmix', type=float, default=0,
|
||||
help='cutmix alpha, cutmix enabled if > 0.')
|
||||
parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
|
||||
help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
|
||||
parser.add_argument('--mixup_prob', type=float, default=1.0,
|
||||
help='Probability of performing mixup or cutmix when either/both is enabled')
|
||||
parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
|
||||
help='Probability of switching to cutmix when both mixup and cutmix enabled')
|
||||
parser.add_argument('--mixup_mode', type=str, default='batch',
|
||||
help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
|
||||
|
||||
# * Finetuning params
|
||||
parser.add_argument('--finetune', default='',
|
||||
help='finetune from checkpoint')
|
||||
parser.add_argument('--model_key', default='model|module', type=str)
|
||||
parser.add_argument('--model_prefix', default='', type=str)
|
||||
parser.add_argument('--init_scale', default=0.001, type=float)
|
||||
parser.add_argument('--use_mean_pooling', action='store_true')
|
||||
parser.set_defaults(use_mean_pooling=True)
|
||||
parser.add_argument('--use_cls', action='store_false', dest='use_mean_pooling')
|
||||
parser.add_argument('--disable_weight_decay_on_rel_pos_bias', action='store_true', default=False)
|
||||
|
||||
# Dataset parameters
|
||||
parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
|
||||
help='dataset path')
|
||||
parser.add_argument('--eval_data_path', default=None, type=str,
|
||||
help='dataset path for evaluation')
|
||||
parser.add_argument('--nb_classes', default=0, type=int,
|
||||
help='number of the classification types')
|
||||
parser.add_argument('--imagenet_default_mean_and_std', default=False, action='store_true')
|
||||
|
||||
parser.add_argument('--data_set', default='IMNET', choices=['CIFAR', 'IMNET', 'image_folder', "rvlcdip", "rvlcdip_wds"],
|
||||
type=str, help='ImageNet dataset path')
|
||||
parser.add_argument('--output_dir', default='',
|
||||
help='path where to save, empty for no saving')
|
||||
parser.add_argument('--log_dir', default=None,
|
||||
help='path where to tensorboard log')
|
||||
parser.add_argument('--device', default='cuda',
|
||||
help='device to use for training / testing')
|
||||
parser.add_argument('--seed', default=0, type=int)
|
||||
parser.add_argument('--resume', default='',
|
||||
help='resume from checkpoint')
|
||||
parser.add_argument('--auto_resume', action='store_true')
|
||||
parser.add_argument('--no_auto_resume', action='store_false', dest='auto_resume')
|
||||
parser.set_defaults(auto_resume=True)
|
||||
|
||||
parser.add_argument('--save_ckpt', action='store_true')
|
||||
parser.add_argument('--no_save_ckpt', action='store_false', dest='save_ckpt')
|
||||
parser.set_defaults(save_ckpt=True)
|
||||
|
||||
parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
|
||||
help='start epoch')
|
||||
parser.add_argument('--eval', action='store_true',
|
||||
help='Perform evaluation only')
|
||||
parser.add_argument('--dist_eval', action='store_true', default=False,
|
||||
help='Enabling distributed evaluation')
|
||||
parser.add_argument('--num_workers', default=10, type=int)
|
||||
parser.add_argument('--pin_mem', action='store_true',
|
||||
help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
|
||||
parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
|
||||
parser.set_defaults(pin_mem=True)
|
||||
|
||||
# distributed training parameters
|
||||
parser.add_argument('--world_size', default=1, type=int,
|
||||
help='number of distributed processes')
|
||||
parser.add_argument('--local_rank', default=-1, type=int)
|
||||
parser.add_argument('--dist_on_itp', action='store_true')
|
||||
parser.add_argument('--dist_url', default='env://',
|
||||
help='url used to set up distributed training')
|
||||
|
||||
parser.add_argument('--enable_deepspeed', action='store_true', default=False)
|
||||
parser.add_argument('--zero_stage', default=0, type=int,
|
||||
help='ZeRO optimizer stage (default: 0)')
|
||||
|
||||
known_args, _ = parser.parse_known_args()
|
||||
|
||||
if known_args.enable_deepspeed:
|
||||
try:
|
||||
import deepspeed
|
||||
from deepspeed import DeepSpeedConfig
|
||||
parser = deepspeed.add_config_arguments(parser)
|
||||
ds_init = deepspeed.initialize
|
||||
except:
|
||||
print("Please 'pip install deepspeed==0.4.0'")
|
||||
exit(0)
|
||||
else:
|
||||
ds_init = None
|
||||
|
||||
return parser.parse_args(), ds_init
|
||||
|
||||
|
||||
def main(args, ds_init):
|
||||
utils.init_distributed_mode(args)
|
||||
|
||||
if ds_init is not None:
|
||||
utils.create_ds_config(args)
|
||||
|
||||
print(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)
|
||||
# random.seed(seed)
|
||||
|
||||
cudnn.benchmark = True
|
||||
|
||||
dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
|
||||
if args.disable_eval_during_finetuning:
|
||||
dataset_val = None
|
||||
else:
|
||||
dataset_val, _ = build_dataset(is_train=False, args=args)
|
||||
|
||||
if True: # args.distributed:
|
||||
num_tasks = utils.get_world_size()
|
||||
global_rank = utils.get_rank()
|
||||
if not isinstance(dataset_train, torch.utils.data.IterableDataset):
|
||||
sampler_train = torch.utils.data.DistributedSampler(
|
||||
dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
|
||||
)
|
||||
print("Sampler_train = %s" % str(sampler_train))
|
||||
if args.dist_eval:
|
||||
if len(dataset_val) % num_tasks != 0:
|
||||
print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
|
||||
'This will slightly alter validation results as extra duplicate entries are added to achieve '
|
||||
'equal num of samples per-process.')
|
||||
sampler_val = torch.utils.data.DistributedSampler(
|
||||
dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
|
||||
else:
|
||||
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
|
||||
else:
|
||||
sampler_train = torch.utils.data.RandomSampler(dataset_train)
|
||||
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
|
||||
|
||||
if 'AMLT_OUTPUT_DIR' in os.environ:
|
||||
args.log_dir = os.environ['AMLT_OUTPUT_DIR']
|
||||
print(f'update log_dir to {args.log_dir}')
|
||||
if global_rank == 0 and args.log_dir is not None:
|
||||
os.makedirs(args.log_dir, exist_ok=True)
|
||||
log_writer = utils.TensorboardLogger(log_dir=args.log_dir)
|
||||
else:
|
||||
log_writer = None
|
||||
|
||||
dataset_size_train = len(dataset_train)
|
||||
if isinstance(dataset_train, torch.utils.data.IterableDataset):
|
||||
dataset_train = dataset_train.batched(args.batch_size, partial=False)
|
||||
data_loader_train = wds.WebLoader(
|
||||
dataset_train, num_workers=args.num_workers, batch_size=None, shuffle=False, )
|
||||
data_loader_train = data_loader_train.ddp_equalize(dataset_size_train // args.batch_size, with_length=True)
|
||||
else:
|
||||
data_loader_train = torch.utils.data.DataLoader(
|
||||
dataset_train, sampler=sampler_train,
|
||||
batch_size=args.batch_size,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=args.pin_mem,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
|
||||
if dataset_val is not None:
|
||||
dataset_size_val = len(dataset_val)
|
||||
if not isinstance(dataset_val, torch.utils.data.IterableDataset):
|
||||
data_loader_val = torch.utils.data.DataLoader(
|
||||
dataset_val, sampler=sampler_val,
|
||||
batch_size=int(1.5 * args.batch_size),
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=args.pin_mem,
|
||||
drop_last=False
|
||||
)
|
||||
else:
|
||||
dataset_val = dataset_val.batched(args.batch_size, partial=False)
|
||||
data_loader_val = wds.WebLoader(
|
||||
dataset_val, num_workers=args.num_workers, batch_size=None, shuffle=False, )
|
||||
data_loader_val = data_loader_val.ddp_equalize(dataset_size_val // args.batch_size, with_length=True)
|
||||
else:
|
||||
data_loader_val = None
|
||||
|
||||
mixup_fn = None
|
||||
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
|
||||
if mixup_active:
|
||||
print("Mixup is activated!")
|
||||
mixup_fn = Mixup(
|
||||
mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
|
||||
prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
|
||||
label_smoothing=args.smoothing, num_classes=args.nb_classes)
|
||||
|
||||
if "beit" not in args.model:
|
||||
model = create_model(args.model, pretrained=False, num_classes=args.nb_classes, distilled=False)
|
||||
else:
|
||||
model = create_model(
|
||||
args.model,
|
||||
pretrained=False,
|
||||
num_classes=args.nb_classes,
|
||||
drop_rate=args.drop,
|
||||
drop_path_rate=args.drop_path,
|
||||
attn_drop_rate=args.attn_drop_rate,
|
||||
drop_block_rate=None,
|
||||
use_mean_pooling=args.use_mean_pooling,
|
||||
init_scale=args.init_scale,
|
||||
use_rel_pos_bias=args.rel_pos_bias,
|
||||
use_abs_pos_emb=args.abs_pos_emb,
|
||||
init_values=args.layer_scale_init_value,
|
||||
)
|
||||
|
||||
patch_size = model.patch_embed.patch_size
|
||||
print("Patch size = %s" % str(patch_size))
|
||||
args.window_size = (args.input_size // patch_size[0], args.input_size // patch_size[1])
|
||||
args.patch_size = patch_size
|
||||
|
||||
if args.finetune:
|
||||
if args.finetune.startswith('https'):
|
||||
checkpoint = torch.hub.load_state_dict_from_url(
|
||||
args.finetune, map_location='cpu', check_hash=False)
|
||||
else:
|
||||
checkpoint = torch.load(args.finetune, map_location='cpu')
|
||||
|
||||
print("Load ckpt from %s" % args.finetune)
|
||||
checkpoint_model = None
|
||||
for model_key in args.model_key.split('|'):
|
||||
if model_key in checkpoint:
|
||||
checkpoint_model = checkpoint[model_key]
|
||||
print("Load state_dict by model_key = %s" % model_key)
|
||||
break
|
||||
if checkpoint_model is None:
|
||||
checkpoint_model = checkpoint
|
||||
state_dict = model.state_dict()
|
||||
for k in ['head.weight', 'head.bias']:
|
||||
if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
|
||||
print(f"Removing key {k} from pretrained checkpoint")
|
||||
del checkpoint_model[k]
|
||||
|
||||
if getattr(model, "use_rel_pos_bias", False) and "rel_pos_bias.relative_position_bias_table" in checkpoint_model:
|
||||
print("Expand the shared relative position embedding to each transformer block. ")
|
||||
num_layers = model.get_num_layers()
|
||||
rel_pos_bias = checkpoint_model["rel_pos_bias.relative_position_bias_table"]
|
||||
for i in range(num_layers):
|
||||
checkpoint_model["blocks.%d.attn.relative_position_bias_table" % i] = rel_pos_bias.clone()
|
||||
|
||||
checkpoint_model.pop("rel_pos_bias.relative_position_bias_table")
|
||||
|
||||
all_keys = list(checkpoint_model.keys())
|
||||
for key in all_keys:
|
||||
if "relative_position_index" in key:
|
||||
checkpoint_model.pop(key)
|
||||
|
||||
if "relative_position_bias_table" in key:
|
||||
rel_pos_bias = checkpoint_model[key]
|
||||
src_num_pos, num_attn_heads = rel_pos_bias.size()
|
||||
dst_num_pos, _ = model.state_dict()[key].size()
|
||||
dst_patch_shape = model.patch_embed.patch_shape
|
||||
if dst_patch_shape[0] != dst_patch_shape[1]:
|
||||
raise NotImplementedError()
|
||||
num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
|
||||
src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
|
||||
dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
|
||||
if src_size != dst_size:
|
||||
print("Position interpolate for %s from %dx%d to %dx%d" % (
|
||||
key, src_size, src_size, dst_size, dst_size))
|
||||
extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
|
||||
rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
|
||||
|
||||
def geometric_progression(a, r, n):
|
||||
return a * (1.0 - r ** n) / (1.0 - r)
|
||||
|
||||
left, right = 1.01, 1.5
|
||||
while right - left > 1e-6:
|
||||
q = (left + right) / 2.0
|
||||
gp = geometric_progression(1, q, src_size // 2)
|
||||
if gp > dst_size // 2:
|
||||
right = q
|
||||
else:
|
||||
left = q
|
||||
|
||||
# if q > 1.090307:
|
||||
# q = 1.090307
|
||||
|
||||
dis = []
|
||||
cur = 1
|
||||
for i in range(src_size // 2):
|
||||
dis.append(cur)
|
||||
cur += q ** (i + 1)
|
||||
|
||||
r_ids = [-_ for _ in reversed(dis)]
|
||||
|
||||
x = r_ids + [0] + dis
|
||||
y = r_ids + [0] + dis
|
||||
|
||||
t = dst_size // 2.0
|
||||
dx = np.arange(-t, t + 0.1, 1.0)
|
||||
dy = np.arange(-t, t + 0.1, 1.0)
|
||||
|
||||
print("Original positions = %s" % str(x))
|
||||
print("Target positions = %s" % str(dx))
|
||||
|
||||
all_rel_pos_bias = []
|
||||
|
||||
for i in range(num_attn_heads):
|
||||
z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
|
||||
f = interpolate.interp2d(x, y, z, kind='cubic')
|
||||
all_rel_pos_bias.append(
|
||||
torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
|
||||
|
||||
rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
|
||||
|
||||
new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
|
||||
checkpoint_model[key] = new_rel_pos_bias
|
||||
|
||||
# interpolate position embedding
|
||||
if 'pos_embed' in checkpoint_model:
|
||||
pos_embed_checkpoint = checkpoint_model['pos_embed']
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.patch_embed.num_patches
|
||||
num_extra_tokens = model.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
new_size = int(num_patches ** 0.5)
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
checkpoint_model['pos_embed'] = new_pos_embed
|
||||
|
||||
utils.load_state_dict(model, checkpoint_model, prefix=args.model_prefix)
|
||||
# model.load_state_dict(checkpoint_model, strict=False)
|
||||
|
||||
model.to(device)
|
||||
|
||||
model_ema = None
|
||||
if args.model_ema:
|
||||
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
|
||||
model_ema = ModelEma(
|
||||
model,
|
||||
decay=args.model_ema_decay,
|
||||
device='cpu' if args.model_ema_force_cpu else '',
|
||||
resume='')
|
||||
print("Using EMA with decay = %.8f" % args.model_ema_decay)
|
||||
|
||||
model_without_ddp = model
|
||||
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
print("Model = %s" % str(model_without_ddp))
|
||||
print('number of params:', n_parameters)
|
||||
|
||||
total_batch_size = args.batch_size * args.update_freq * utils.get_world_size()
|
||||
num_training_steps_per_epoch = dataset_size_train // total_batch_size
|
||||
print("LR = %.8f" % args.lr)
|
||||
print("Batch size = %d" % total_batch_size)
|
||||
print("Update frequent = %d" % args.update_freq)
|
||||
print("Number of training examples = %d" % dataset_size_train)
|
||||
print("Number of training training per epoch = %d" % num_training_steps_per_epoch)
|
||||
|
||||
# num_layers = model_without_ddp.get_num_layers()
|
||||
num_layers = len(model_without_ddp.blocks)
|
||||
if args.layer_decay < 1.0:
|
||||
assigner = LayerDecayValueAssigner(list(args.layer_decay ** (num_layers + 1 - i) for i in range(num_layers + 2)))
|
||||
else:
|
||||
assigner = None
|
||||
|
||||
if assigner is not None:
|
||||
print("Assigned values = %s" % str(assigner.values))
|
||||
|
||||
skip_weight_decay_list = model.no_weight_decay()
|
||||
if args.disable_weight_decay_on_rel_pos_bias:
|
||||
for i in range(num_layers):
|
||||
skip_weight_decay_list.add("blocks.%d.attn.relative_position_bias_table" % i)
|
||||
|
||||
if args.distributed:
|
||||
torch.distributed.barrier()
|
||||
if args.enable_deepspeed:
|
||||
loss_scaler = None
|
||||
optimizer_params = get_parameter_groups(
|
||||
model, args.weight_decay, skip_weight_decay_list,
|
||||
assigner.get_layer_id if assigner is not None else None,
|
||||
assigner.get_scale if assigner is not None else None)
|
||||
model, optimizer, _, _ = ds_init(
|
||||
args=args, model=model, model_parameters=optimizer_params,
|
||||
dist_init_required=not args.distributed,
|
||||
)
|
||||
|
||||
print("model.gradient_accumulation_steps() = %d" % model.gradient_accumulation_steps())
|
||||
assert model.gradient_accumulation_steps() == args.update_freq
|
||||
else:
|
||||
if args.distributed:
|
||||
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
|
||||
model_without_ddp = model.module
|
||||
|
||||
optimizer = create_optimizer(
|
||||
args, model_without_ddp, skip_list=skip_weight_decay_list,
|
||||
get_num_layer=assigner.get_layer_id if assigner is not None else None,
|
||||
get_layer_scale=assigner.get_scale if assigner is not None else None)
|
||||
loss_scaler = NativeScaler()
|
||||
|
||||
print("Use step level LR scheduler!")
|
||||
lr_schedule_values = utils.cosine_scheduler(
|
||||
args.lr, args.min_lr, args.epochs, num_training_steps_per_epoch,
|
||||
warmup_epochs=args.warmup_epochs, warmup_steps=args.warmup_steps,
|
||||
)
|
||||
if args.weight_decay_end is None:
|
||||
args.weight_decay_end = args.weight_decay
|
||||
wd_schedule_values = utils.cosine_scheduler(
|
||||
args.weight_decay, args.weight_decay_end, args.epochs, num_training_steps_per_epoch)
|
||||
print("Max WD = %.7f, Min WD = %.7f" % (max(wd_schedule_values), min(wd_schedule_values)))
|
||||
|
||||
if mixup_fn is not None:
|
||||
# smoothing is handled with mixup label transform
|
||||
criterion = SoftTargetCrossEntropy()
|
||||
elif args.smoothing > 0.:
|
||||
criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
|
||||
else:
|
||||
criterion = torch.nn.CrossEntropyLoss()
|
||||
|
||||
print("criterion = %s" % str(criterion))
|
||||
|
||||
utils.auto_load_model(
|
||||
args=args, model=model, model_without_ddp=model_without_ddp,
|
||||
optimizer=optimizer, loss_scaler=loss_scaler, model_ema=model_ema)
|
||||
|
||||
if args.eval:
|
||||
test_stats = evaluate(data_loader_val, model, device)
|
||||
print(f"Accuracy of the network on the {dataset_size_val} test images: {test_stats['acc1']:.1f}%")
|
||||
exit(0)
|
||||
|
||||
print(f"Start training for {args.epochs} epochs")
|
||||
start_time = time.time()
|
||||
max_accuracy = 0.0
|
||||
for epoch in range(args.start_epoch, args.epochs):
|
||||
if args.distributed:
|
||||
sampler = getattr(data_loader_train, "sampler", None)
|
||||
if sampler is not None and hasattr(sampler, "set_epoch"):
|
||||
sampler.set_epoch(epoch)
|
||||
if log_writer is not None:
|
||||
log_writer.set_step(epoch * num_training_steps_per_epoch * args.update_freq)
|
||||
train_stats = train_one_epoch(
|
||||
model, criterion, data_loader_train, optimizer,
|
||||
device, epoch, loss_scaler, args.clip_grad, model_ema, mixup_fn,
|
||||
log_writer=log_writer, start_steps=epoch * num_training_steps_per_epoch,
|
||||
lr_schedule_values=lr_schedule_values, wd_schedule_values=wd_schedule_values,
|
||||
num_training_steps_per_epoch=num_training_steps_per_epoch, update_freq=args.update_freq,
|
||||
)
|
||||
if args.output_dir and args.save_ckpt:
|
||||
if (epoch + 1) % args.save_ckpt_freq == 0 or epoch + 1 == args.epochs:
|
||||
utils.save_model(
|
||||
args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
|
||||
loss_scaler=loss_scaler, epoch=epoch, model_ema=model_ema)
|
||||
if data_loader_val is not None and ((epoch + 1) % args.eval_freq == 0 or epoch + 1 == args.epochs):
|
||||
test_stats = evaluate(data_loader_val, model, device)
|
||||
print(f"Accuracy of the network on the {dataset_size_val} test images: {test_stats['acc1']:.1f}%")
|
||||
if max_accuracy < test_stats["acc1"]:
|
||||
max_accuracy = test_stats["acc1"]
|
||||
if args.output_dir and args.save_ckpt:
|
||||
utils.save_model(
|
||||
args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
|
||||
loss_scaler=loss_scaler, epoch="best", model_ema=model_ema)
|
||||
|
||||
print(f'Max accuracy: {max_accuracy:.2f}%')
|
||||
if log_writer is not None:
|
||||
log_writer.update(test_acc1=test_stats['acc1'], head="perf", step=epoch)
|
||||
log_writer.update(test_acc5=test_stats['acc5'], head="perf", step=epoch)
|
||||
log_writer.update(test_loss=test_stats['loss'], head="perf", step=epoch)
|
||||
|
||||
log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
|
||||
**{f'test_{k}': v for k, v in test_stats.items()},
|
||||
'epoch': epoch,
|
||||
'n_parameters': n_parameters}
|
||||
else:
|
||||
log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
|
||||
# **{f'test_{k}': v for k, v in test_stats.items()},
|
||||
'epoch': epoch,
|
||||
'n_parameters': n_parameters}
|
||||
|
||||
if args.output_dir and utils.is_main_process():
|
||||
if log_writer is not None:
|
||||
log_writer.flush()
|
||||
with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(log_stats) + "\n")
|
||||
|
||||
total_time = time.time() - start_time
|
||||
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
|
||||
print('Training time {}'.format(total_time_str))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
opts, ds_init = get_args()
|
||||
if opts.output_dir:
|
||||
Path(opts.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
main(opts, ds_init)
|
||||
@@ -0,0 +1,132 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Based on timm code bases
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# --------------------------------------------------------'
|
||||
import math
|
||||
import random
|
||||
import warnings
|
||||
|
||||
import torchvision.transforms.functional as F
|
||||
from timm.data.transforms import interp_mode_to_str, _RANDOM_INTERPOLATION, str_to_interp_mode
|
||||
|
||||
|
||||
class RandomResizedCropAndInterpolationWithTwoPic:
|
||||
"""Crop the given PIL Image to random size and aspect ratio with random interpolation.
|
||||
|
||||
A crop of random size (default: of 0.08 to 1.0) of the original size and a random
|
||||
aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
|
||||
is finally resized to given size.
|
||||
This is popularly used to train the Inception networks.
|
||||
|
||||
Args:
|
||||
size: expected output size of each edge
|
||||
scale: range of size of the origin size cropped
|
||||
ratio: range of aspect ratio of the origin aspect ratio cropped
|
||||
interpolation: Default: PIL.Image.BILINEAR
|
||||
"""
|
||||
|
||||
def __init__(self, size, second_size=None, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.),
|
||||
interpolation='bilinear', second_interpolation='lanczos'):
|
||||
if isinstance(size, tuple):
|
||||
self.size = size
|
||||
else:
|
||||
self.size = (size, size)
|
||||
if second_size is not None:
|
||||
if isinstance(second_size, tuple):
|
||||
self.second_size = second_size
|
||||
else:
|
||||
self.second_size = (second_size, second_size)
|
||||
else:
|
||||
self.second_size = None
|
||||
if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
|
||||
warnings.warn("range should be of kind (min, max)")
|
||||
|
||||
if interpolation == 'random':
|
||||
self.interpolation = _RANDOM_INTERPOLATION
|
||||
else:
|
||||
self.interpolation = str_to_interp_mode(interpolation)
|
||||
self.second_interpolation = str_to_interp_mode(second_interpolation)
|
||||
self.scale = scale
|
||||
self.ratio = ratio
|
||||
|
||||
@staticmethod
|
||||
def get_params(img, scale, ratio):
|
||||
"""Get parameters for ``crop`` for a random sized crop.
|
||||
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped.
|
||||
scale (tuple): range of size of the origin size cropped
|
||||
ratio (tuple): range of aspect ratio of the origin aspect ratio cropped
|
||||
|
||||
Returns:
|
||||
tuple: params (i, j, h, w) to be passed to ``crop`` for a random
|
||||
sized crop.
|
||||
"""
|
||||
area = img.size[0] * img.size[1]
|
||||
|
||||
for attempt in range(10):
|
||||
target_area = random.uniform(*scale) * area
|
||||
log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
|
||||
aspect_ratio = math.exp(random.uniform(*log_ratio))
|
||||
|
||||
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
|
||||
if w <= img.size[0] and h <= img.size[1]:
|
||||
i = random.randint(0, img.size[1] - h)
|
||||
j = random.randint(0, img.size[0] - w)
|
||||
return i, j, h, w
|
||||
|
||||
# Fallback to central crop
|
||||
in_ratio = img.size[0] / img.size[1]
|
||||
if in_ratio < min(ratio):
|
||||
w = img.size[0]
|
||||
h = int(round(w / min(ratio)))
|
||||
elif in_ratio > max(ratio):
|
||||
h = img.size[1]
|
||||
w = int(round(h * max(ratio)))
|
||||
else: # whole image
|
||||
w = img.size[0]
|
||||
h = img.size[1]
|
||||
i = (img.size[1] - h) // 2
|
||||
j = (img.size[0] - w) // 2
|
||||
return i, j, h, w
|
||||
|
||||
def __call__(self, img):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped and resized.
|
||||
|
||||
Returns:
|
||||
PIL Image: Randomly cropped and resized image.
|
||||
"""
|
||||
i, j, h, w = self.get_params(img, self.scale, self.ratio)
|
||||
if isinstance(self.interpolation, (tuple, list)):
|
||||
interpolation = random.choice(self.interpolation)
|
||||
else:
|
||||
interpolation = self.interpolation
|
||||
if self.second_size is None:
|
||||
return F.resized_crop(img, i, j, h, w, self.size, interpolation)
|
||||
else:
|
||||
return F.resized_crop(img, i, j, h, w, self.size, interpolation), \
|
||||
F.resized_crop(img, i, j, h, w, self.second_size, self.second_interpolation)
|
||||
|
||||
def __repr__(self):
|
||||
if isinstance(self.interpolation, (tuple, list)):
|
||||
interpolate_str = ' '.join([interp_mode_to_str(x) for x in self.interpolation])
|
||||
else:
|
||||
interpolate_str = interp_mode_to_str(self.interpolation)
|
||||
format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
|
||||
format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
|
||||
format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
|
||||
format_string += ', interpolation={0}'.format(interpolate_str)
|
||||
if self.second_size is not None:
|
||||
format_string += ', second_size={0}'.format(self.second_size)
|
||||
format_string += ', second_interpolation={0}'.format(interp_mode_to_str(self.second_interpolation))
|
||||
format_string += ')'
|
||||
return format_string
|
||||
@@ -0,0 +1,523 @@
|
||||
# --------------------------------------------------------
|
||||
# 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 Hangbo Bao
|
||||
# Based on timm, DINO and DeiT code bases
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# https://github.com/facebookresearch/deit
|
||||
# https://github.com/facebookresearch/dino
|
||||
# --------------------------------------------------------'
|
||||
import datetime
|
||||
import io
|
||||
import os
|
||||
import math
|
||||
import time
|
||||
import json
|
||||
from collections import defaultdict, deque
|
||||
import datetime
|
||||
import numpy as np
|
||||
from timm.utils import get_state_dict
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch._six import inf
|
||||
# from modeling_discrete_vae import Dalle_VAE, DiscreteVAE, DiscreteVAE2, VQGanVAE, DiscreteVAEforBEiT
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
class SmoothedValue(object):
|
||||
"""Track a series of values and provide access to smoothed values over a
|
||||
window or the global series average.
|
||||
"""
|
||||
|
||||
def __init__(self, window_size=20, fmt=None):
|
||||
if fmt is None:
|
||||
fmt = "{median:.4f} ({global_avg:.4f})"
|
||||
self.deque = deque(maxlen=window_size)
|
||||
self.total = 0.0
|
||||
self.count = 0
|
||||
self.fmt = fmt
|
||||
|
||||
def update(self, value, n=1):
|
||||
self.deque.append(value)
|
||||
self.count += n
|
||||
self.total += value * n
|
||||
|
||||
def synchronize_between_processes(self):
|
||||
"""
|
||||
Warning: does not synchronize the deque!
|
||||
"""
|
||||
if not is_dist_avail_and_initialized():
|
||||
return
|
||||
t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
|
||||
dist.barrier()
|
||||
dist.all_reduce(t)
|
||||
t = t.tolist()
|
||||
self.count = int(t[0])
|
||||
self.total = t[1]
|
||||
|
||||
@property
|
||||
def median(self):
|
||||
d = torch.tensor(list(self.deque))
|
||||
return d.median().item()
|
||||
|
||||
@property
|
||||
def avg(self):
|
||||
d = torch.tensor(list(self.deque), dtype=torch.float32)
|
||||
return d.mean().item()
|
||||
|
||||
@property
|
||||
def global_avg(self):
|
||||
return self.total / self.count
|
||||
|
||||
@property
|
||||
def max(self):
|
||||
return max(self.deque)
|
||||
|
||||
@property
|
||||
def value(self):
|
||||
return self.deque[-1]
|
||||
|
||||
def __str__(self):
|
||||
return self.fmt.format(
|
||||
median=self.median,
|
||||
avg=self.avg,
|
||||
global_avg=self.global_avg,
|
||||
max=self.max,
|
||||
value=self.value)
|
||||
|
||||
|
||||
class MetricLogger(object):
|
||||
def __init__(self, delimiter="\t"):
|
||||
self.meters = defaultdict(SmoothedValue)
|
||||
self.delimiter = delimiter
|
||||
|
||||
def update(self, **kwargs):
|
||||
for k, v in kwargs.items():
|
||||
if v is None:
|
||||
continue
|
||||
if isinstance(v, torch.Tensor):
|
||||
v = v.item()
|
||||
assert isinstance(v, (float, int))
|
||||
self.meters[k].update(v)
|
||||
|
||||
def __getattr__(self, attr):
|
||||
if attr in self.meters:
|
||||
return self.meters[attr]
|
||||
if attr in self.__dict__:
|
||||
return self.__dict__[attr]
|
||||
raise AttributeError("'{}' object has no attribute '{}'".format(
|
||||
type(self).__name__, attr))
|
||||
|
||||
def __str__(self):
|
||||
loss_str = []
|
||||
for name, meter in self.meters.items():
|
||||
loss_str.append(
|
||||
"{}: {}".format(name, str(meter))
|
||||
)
|
||||
return self.delimiter.join(loss_str)
|
||||
|
||||
def synchronize_between_processes(self):
|
||||
for meter in self.meters.values():
|
||||
meter.synchronize_between_processes()
|
||||
|
||||
def add_meter(self, name, meter):
|
||||
self.meters[name] = meter
|
||||
|
||||
def log_every(self, iterable, print_freq, header=None):
|
||||
i = 0
|
||||
if not header:
|
||||
header = ''
|
||||
start_time = time.time()
|
||||
end = time.time()
|
||||
iter_time = SmoothedValue(fmt='{avg:.4f}')
|
||||
data_time = SmoothedValue(fmt='{avg:.4f}')
|
||||
space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
|
||||
log_msg = [
|
||||
header,
|
||||
'[{0' + space_fmt + '}/{1}]',
|
||||
'eta: {eta}',
|
||||
'{meters}',
|
||||
'time: {time}',
|
||||
'data: {data}'
|
||||
]
|
||||
if torch.cuda.is_available():
|
||||
log_msg.append('max mem: {memory:.0f}')
|
||||
log_msg = self.delimiter.join(log_msg)
|
||||
MB = 1024.0 * 1024.0
|
||||
for obj in iterable:
|
||||
data_time.update(time.time() - end)
|
||||
yield obj
|
||||
iter_time.update(time.time() - end)
|
||||
if i % print_freq == 0 or i == len(iterable) - 1:
|
||||
eta_seconds = iter_time.global_avg * (len(iterable) - i)
|
||||
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
|
||||
if torch.cuda.is_available():
|
||||
print(log_msg.format(
|
||||
i, len(iterable), eta=eta_string,
|
||||
meters=str(self),
|
||||
time=str(iter_time), data=str(data_time),
|
||||
memory=torch.cuda.max_memory_allocated() / MB))
|
||||
else:
|
||||
print(log_msg.format(
|
||||
i, len(iterable), eta=eta_string,
|
||||
meters=str(self),
|
||||
time=str(iter_time), data=str(data_time)))
|
||||
i += 1
|
||||
end = time.time()
|
||||
total_time = time.time() - start_time
|
||||
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
|
||||
print('{} Total time: {} ({:.4f} s / it)'.format(
|
||||
header, total_time_str, total_time / len(iterable)))
|
||||
|
||||
|
||||
class TensorboardLogger(object):
|
||||
def __init__(self, log_dir):
|
||||
self.writer = SummaryWriter(log_dir=log_dir)
|
||||
self.step = 0
|
||||
|
||||
def set_step(self, step=None):
|
||||
if step is not None:
|
||||
self.step = step
|
||||
else:
|
||||
self.step += 1
|
||||
|
||||
def update(self, head='scalar', step=None, **kwargs):
|
||||
for k, v in kwargs.items():
|
||||
if v is None:
|
||||
continue
|
||||
if isinstance(v, torch.Tensor):
|
||||
v = v.item()
|
||||
assert isinstance(v, (float, int))
|
||||
self.writer.add_scalar(head + "/" + k, v, self.step if step is None else step)
|
||||
|
||||
def flush(self):
|
||||
self.writer.flush()
|
||||
|
||||
|
||||
def _load_checkpoint_for_ema(model_ema, checkpoint):
|
||||
"""
|
||||
Workaround for ModelEma._load_checkpoint to accept an already-loaded object
|
||||
"""
|
||||
mem_file = io.BytesIO()
|
||||
torch.save(checkpoint, mem_file)
|
||||
mem_file.seek(0)
|
||||
model_ema._load_checkpoint(mem_file)
|
||||
|
||||
|
||||
def setup_for_distributed(is_master):
|
||||
"""
|
||||
This function disables printing when not in master process
|
||||
"""
|
||||
import builtins as __builtin__
|
||||
builtin_print = __builtin__.print
|
||||
|
||||
def print(*args, **kwargs):
|
||||
force = kwargs.pop('force', False)
|
||||
if is_master or force:
|
||||
builtin_print(*args, **kwargs)
|
||||
|
||||
__builtin__.print = print
|
||||
|
||||
|
||||
def is_dist_avail_and_initialized():
|
||||
if not dist.is_available():
|
||||
return False
|
||||
if not dist.is_initialized():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_world_size():
|
||||
if not is_dist_avail_and_initialized():
|
||||
return 1
|
||||
return dist.get_world_size()
|
||||
|
||||
|
||||
def get_rank():
|
||||
if not is_dist_avail_and_initialized():
|
||||
return 0
|
||||
return dist.get_rank()
|
||||
|
||||
|
||||
def is_main_process():
|
||||
return get_rank() == 0
|
||||
|
||||
|
||||
def save_on_master(*args, **kwargs):
|
||||
if is_main_process():
|
||||
torch.save(*args, **kwargs)
|
||||
|
||||
|
||||
def init_distributed_mode(args):
|
||||
if args.dist_on_itp:
|
||||
args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
|
||||
args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
|
||||
args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
|
||||
args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
|
||||
os.environ['LOCAL_RANK'] = str(args.gpu)
|
||||
os.environ['RANK'] = str(args.rank)
|
||||
os.environ['WORLD_SIZE'] = str(args.world_size)
|
||||
# ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
|
||||
elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
|
||||
args.rank = int(os.environ["RANK"])
|
||||
args.world_size = int(os.environ['WORLD_SIZE'])
|
||||
args.gpu = int(os.environ['LOCAL_RANK'])
|
||||
elif 'SLURM_PROCID' in os.environ:
|
||||
args.rank = int(os.environ['SLURM_PROCID'])
|
||||
args.gpu = args.rank % torch.cuda.device_count()
|
||||
else:
|
||||
print('Not using distributed mode')
|
||||
args.distributed = False
|
||||
return
|
||||
|
||||
args.distributed = True
|
||||
|
||||
torch.cuda.set_device(args.gpu)
|
||||
args.dist_backend = 'nccl'
|
||||
print('| distributed init (rank {}): {}, gpu {}'.format(
|
||||
args.rank, args.dist_url, args.gpu), flush=True)
|
||||
torch.distributed.init_process_group(
|
||||
backend=args.dist_backend, init_method=args.dist_url,
|
||||
world_size=args.world_size, rank=args.rank,
|
||||
timeout=datetime.timedelta(0, 7200)
|
||||
)
|
||||
torch.distributed.barrier()
|
||||
setup_for_distributed(args.rank == 0)
|
||||
|
||||
|
||||
def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"):
|
||||
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 + '.')
|
||||
|
||||
load(model, prefix=prefix)
|
||||
|
||||
warn_missing_keys = []
|
||||
ignore_missing_keys = []
|
||||
for key in missing_keys:
|
||||
keep_flag = True
|
||||
for ignore_key in ignore_missing.split('|'):
|
||||
if ignore_key in key:
|
||||
keep_flag = False
|
||||
break
|
||||
if keep_flag:
|
||||
warn_missing_keys.append(key)
|
||||
else:
|
||||
ignore_missing_keys.append(key)
|
||||
|
||||
missing_keys = warn_missing_keys
|
||||
|
||||
if len(missing_keys) > 0:
|
||||
print("Weights of {} not initialized from pretrained model: {}".format(
|
||||
model.__class__.__name__, missing_keys))
|
||||
if len(unexpected_keys) > 0:
|
||||
print("Weights from pretrained model not used in {}: {}".format(
|
||||
model.__class__.__name__, unexpected_keys))
|
||||
if len(ignore_missing_keys) > 0:
|
||||
print("Ignored weights of {} not initialized from pretrained model: {}".format(
|
||||
model.__class__.__name__, ignore_missing_keys))
|
||||
if len(error_msgs) > 0:
|
||||
print('\n'.join(error_msgs))
|
||||
|
||||
|
||||
class NativeScalerWithGradNormCount:
|
||||
state_dict_key = "amp_scaler"
|
||||
|
||||
def __init__(self):
|
||||
self._scaler = torch.cuda.amp.GradScaler()
|
||||
|
||||
def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
|
||||
self._scaler.scale(loss).backward(create_graph=create_graph)
|
||||
if update_grad:
|
||||
if clip_grad is not None:
|
||||
assert parameters is not None
|
||||
self._scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place
|
||||
norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
|
||||
else:
|
||||
self._scaler.unscale_(optimizer)
|
||||
norm = get_grad_norm_(parameters)
|
||||
self._scaler.step(optimizer)
|
||||
self._scaler.update()
|
||||
else:
|
||||
norm = None
|
||||
return norm
|
||||
|
||||
def state_dict(self):
|
||||
return self._scaler.state_dict()
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
self._scaler.load_state_dict(state_dict)
|
||||
|
||||
|
||||
def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
|
||||
if isinstance(parameters, torch.Tensor):
|
||||
parameters = [parameters]
|
||||
parameters = [p for p in parameters if p.grad is not None]
|
||||
norm_type = float(norm_type)
|
||||
if len(parameters) == 0:
|
||||
return torch.tensor(0.)
|
||||
device = parameters[0].grad.device
|
||||
if norm_type == inf:
|
||||
total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
|
||||
else:
|
||||
total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
|
||||
return total_norm
|
||||
|
||||
|
||||
def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0,
|
||||
start_warmup_value=0, warmup_steps=-1):
|
||||
warmup_schedule = np.array([])
|
||||
warmup_iters = warmup_epochs * niter_per_ep
|
||||
if warmup_steps > 0:
|
||||
warmup_iters = warmup_steps
|
||||
print("Set warmup steps = %d" % warmup_iters)
|
||||
if warmup_epochs > 0:
|
||||
warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
|
||||
|
||||
iters = np.arange(epochs * niter_per_ep - warmup_iters)
|
||||
schedule = np.array(
|
||||
[final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
|
||||
|
||||
schedule = np.concatenate((warmup_schedule, schedule))
|
||||
|
||||
# assert len(schedule) == epochs * niter_per_ep
|
||||
return schedule
|
||||
|
||||
|
||||
def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
|
||||
output_dir = Path(args.output_dir)
|
||||
epoch_name = str(epoch)
|
||||
if loss_scaler is not None:
|
||||
checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
|
||||
for checkpoint_path in checkpoint_paths:
|
||||
to_save = {
|
||||
'model': model_without_ddp.state_dict(),
|
||||
'optimizer': optimizer.state_dict(),
|
||||
'epoch': epoch,
|
||||
'scaler': loss_scaler.state_dict(),
|
||||
'args': args,
|
||||
}
|
||||
|
||||
if model_ema is not None:
|
||||
to_save['model_ema'] = get_state_dict(model_ema)
|
||||
|
||||
save_on_master(to_save, checkpoint_path)
|
||||
else:
|
||||
client_state = {'epoch': epoch}
|
||||
if model_ema is not None:
|
||||
client_state['model_ema'] = get_state_dict(model_ema)
|
||||
model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
|
||||
|
||||
|
||||
def auto_load_model(args, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
|
||||
output_dir = Path(args.output_dir)
|
||||
if loss_scaler is not None:
|
||||
# torch.amp
|
||||
if args.auto_resume and len(args.resume) == 0:
|
||||
import glob
|
||||
all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*.pth'))
|
||||
latest_ckpt = -1
|
||||
for ckpt in all_checkpoints:
|
||||
t = ckpt.split('-')[-1].split('.')[0]
|
||||
if t.isdigit():
|
||||
latest_ckpt = max(int(t), latest_ckpt)
|
||||
if latest_ckpt >= 0:
|
||||
args.resume = os.path.join(output_dir, 'checkpoint-%d.pth' % latest_ckpt)
|
||||
print("Auto resume checkpoint: %s" % args.resume)
|
||||
|
||||
if args.resume:
|
||||
if args.resume.startswith('https'):
|
||||
checkpoint = torch.hub.load_state_dict_from_url(
|
||||
args.resume, map_location='cpu', check_hash=True)
|
||||
else:
|
||||
checkpoint = torch.load(args.resume, map_location='cpu')
|
||||
model_without_ddp.load_state_dict(checkpoint['model'])
|
||||
print("Resume checkpoint %s" % args.resume)
|
||||
if 'optimizer' in checkpoint and 'epoch' in checkpoint:
|
||||
optimizer.load_state_dict(checkpoint['optimizer'])
|
||||
args.start_epoch = checkpoint['epoch'] + 1
|
||||
if hasattr(args, 'model_ema') and args.model_ema:
|
||||
_load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
|
||||
if 'scaler' in checkpoint:
|
||||
loss_scaler.load_state_dict(checkpoint['scaler'])
|
||||
print("With optim & sched!")
|
||||
else:
|
||||
# deepspeed, only support '--auto_resume'.
|
||||
if args.auto_resume:
|
||||
import glob
|
||||
all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*'))
|
||||
latest_ckpt = -1
|
||||
for ckpt in all_checkpoints:
|
||||
t = ckpt.split('-')[-1].split('.')[0]
|
||||
if t.isdigit():
|
||||
latest_ckpt = max(int(t), latest_ckpt)
|
||||
if latest_ckpt >= 0:
|
||||
args.resume = os.path.join(output_dir, 'checkpoint-%d' % latest_ckpt)
|
||||
print("Auto resume checkpoint: %d" % latest_ckpt)
|
||||
_, client_states = model.load_checkpoint(args.output_dir, tag='checkpoint-%d' % latest_ckpt)
|
||||
args.start_epoch = client_states['epoch'] + 1
|
||||
if model_ema is not None:
|
||||
if args.model_ema:
|
||||
_load_checkpoint_for_ema(model_ema, client_states['model_ema'])
|
||||
|
||||
def create_ds_config(args):
|
||||
args.deepspeed_config = os.path.join(args.output_dir, "deepspeed_config.json")
|
||||
with open(args.deepspeed_config, mode="w") as writer:
|
||||
ds_config = {
|
||||
"train_batch_size": args.batch_size * args.update_freq * get_world_size(),
|
||||
"train_micro_batch_size_per_gpu": args.batch_size,
|
||||
"steps_per_print": 1000,
|
||||
"optimizer": {
|
||||
"type": "Adam",
|
||||
"adam_w_mode": True,
|
||||
"params": {
|
||||
"lr": args.lr,
|
||||
"weight_decay": args.weight_decay,
|
||||
"bias_correction": True,
|
||||
"betas": [
|
||||
0.9,
|
||||
0.999
|
||||
],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"fp16": {
|
||||
"enabled": True,
|
||||
"loss_scale": 0,
|
||||
"initial_scale_power": 16,
|
||||
"loss_scale_window": 1000,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"zero_optimization": {
|
||||
"stage": args.zero_stage
|
||||
},
|
||||
"amp": {
|
||||
"enabled": False,
|
||||
"opt_level": "O2"
|
||||
}
|
||||
}
|
||||
|
||||
if args.clip_grad is not None:
|
||||
ds_config.update({'gradient_clipping': args.clip_grad})
|
||||
|
||||
writer.write(json.dumps(ds_config, indent=2))
|
||||
@@ -0,0 +1,120 @@
|
||||
# DiT for Object Detection
|
||||
|
||||
This folder contains Mask R-CNN Cascade Mask R-CNN running instructions on top of [Detectron2](https://github.com/facebookresearch/detectron2) for PubLayNet and ICDAR 2019 cTDaR.
|
||||
|
||||
## Usage
|
||||
|
||||
### Inference
|
||||
|
||||
The quickest way to try out DiT for document layout analysis is the web demo: [](https://huggingface.co/spaces/nielsr/dit-document-layout-analysis).
|
||||
|
||||
One can run inference using the `inference.py` script. It can be run as follows (from the root of the unilm repository):
|
||||
|
||||
```
|
||||
python ./dit/object_detection/inference.py \
|
||||
--image_path ./dit/object_detection/publaynet_example.jpeg \
|
||||
--output_file_name output.jpg \
|
||||
--config ./dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml \
|
||||
--opts MODEL.WEIGHTS https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth \
|
||||
```
|
||||
|
||||
Make sure that the configuration file (YAML) and PyTorch checkpoint match. The example above uses DiT-base with the Mask R-CNN framework fine-tuned on PubLayNet.
|
||||
|
||||
### Data Preparation
|
||||
|
||||
**PubLayNet**
|
||||
|
||||
Download the data from this [link](https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz?_ga=2.218138265.1825957955.1646384196-1495010506.1633610665) (~96GB). Then extract it to `PATH-to-PubLayNet`.
|
||||
|
||||
A soft link needs to be created to make the data accessible for the program:`ln -s PATH-to-PubLayNet publaynet_data`.
|
||||
|
||||
**ICDAR 2019 cTDaR**
|
||||
|
||||
Download the data from this [link](https://github.com/cndplab-founder/ICDAR2019_cTDaR) (~4GB). Assume path to this repository is named as `PATH-to-ICDARrepo`.
|
||||
|
||||
Then run `python convert_to_coco_format.py --root_dir=PATH-to-ICDARrepo --target_dir=PATH-toICDAR`. Now the path to processed data is `PATH-to-ICDAR`.
|
||||
|
||||
Run the following command to get the adaptively binarized images for archival subset.
|
||||
|
||||
```
|
||||
cp -r PATH-to-ICDAR/trackA_archival PATH-to-ICDAR/at_trackA_archival
|
||||
python adaptive_binarize.py --root_dir PATH-to-ICDAR/at_trackA_archival
|
||||
```
|
||||
|
||||
The binarized archival subset will be in `PATH-to-ICDAR/at_trackA_archival`.
|
||||
|
||||
According to the subset you want to evaluate/fine-tune, a soft link should be created:`ln -s PATH-to-ICDAR/trackA_modern data` or `ln -s PATH-to-ICDAR/at_trackA_archival data`.
|
||||
|
||||
### Evaluation
|
||||
|
||||
Following commands provide two examples to evaluate the fine-tuned checkpoints.
|
||||
|
||||
The config files can be found in `icdar19_configs` and `publaynet_configs`.
|
||||
|
||||
1) Evaluate the fine-tuned checkpoint of DiT-Base with Mask R-CNN on PublayNet:
|
||||
```bash
|
||||
python train_net.py --config-file publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir>
|
||||
```
|
||||
|
||||
2) Evaluate the fine-tuned checkpoint of DiT-Large with Cascade Mask R-CNN on ICDAR 2019 cTDaR archival subset (make sure you have created a soft link from `PATH-to-ICDAR/at_trackA_archival` to `data`):
|
||||
```bash
|
||||
python train_net.py --config-file icdar19_configs/cascade/cascade_dit_large.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir>
|
||||
```
|
||||
|
||||
**Note**: We have fixed the **bug** in the [ICDAR2019 measurement tool](https://github.com/cndplab-founder/ctdar_measurement_tool) during integrating the tool into our code. If you use the tool to get the evaluation score, please modify the [code](https://github.com/cndplab-founder/ctdar_measurement_tool/blob/738456d3164a838ffaeefe7d1b5e64f3a4368a0e/evaluate.py#L146
|
||||
) as follows:
|
||||
```bash
|
||||
...
|
||||
# print(each_file)
|
||||
|
||||
# for file in gt_file_lst:
|
||||
# if file.split(".") != "xml":
|
||||
# gt_file_lst.remove(file)
|
||||
# # print(gt_file_lst)
|
||||
|
||||
# Comment the code above and add the code below
|
||||
for i in range(len(gt_file_lst) - 1, -1, -1):
|
||||
if gt_file_lst[i].split(".")[-1] != "xml":
|
||||
del gt_file_lst[i]
|
||||
|
||||
if len(gt_file_lst) > 0:
|
||||
...
|
||||
```
|
||||
|
||||
### Training
|
||||
The following commands provide two examples to train the Mask R-CNN/Cascade Mask R-CNN with DiT backbone on 8 32GB Nvidia V100 GPUs.
|
||||
|
||||
1) Fine-tune DiT-Base with Cascade Mask R-CNN on PublayNet:
|
||||
```bash
|
||||
python train_net.py --config-file publaynet_configs/cascade/cascade_dit_base.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Base_file_path or link> OUTPUT_DIR <your_output_dir>
|
||||
```
|
||||
|
||||
|
||||
2) Fine-tune DiT-Large with Mask R-CNN on ICDAR 2019 cTDaR modern:
|
||||
```bash
|
||||
python train_net.py --config-file icdar19_configs/markrcnn/maskrcnn_dit_large.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Large_file_path or link> OUTPUT_DIR <your_output_dir>
|
||||
```
|
||||
|
||||
|
||||
|
||||
[Detectron2's document](https://detectron2.readthedocs.io/en/latest/tutorials/getting_started.html) may help you for more details.
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider citing our work:
|
||||
```
|
||||
@misc{li2022dit,
|
||||
title={DiT: Self-supervised Pre-training for Document Image Transformer},
|
||||
author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
|
||||
year={2022},
|
||||
eprint={2203.02378},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
|
||||
## Acknowledgment
|
||||
Thanks to [Detectron2](https://github.com/facebookresearch/detectron2) for Mask R-CNN and Cascade Mask R-CNN implementation.
|
||||
@@ -0,0 +1,33 @@
|
||||
import argparse
|
||||
import os
|
||||
import cv2
|
||||
import tqdm
|
||||
|
||||
|
||||
def convert(fn):
|
||||
# given a file name, convert it into binary and store at the same position
|
||||
img = cv2.imread(fn)
|
||||
gim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
gim = cv2.adaptiveThreshold(gim, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 45, 11)
|
||||
g3im = cv2.cvtColor(gim, cv2.COLOR_GRAY2BGR)
|
||||
cv2.imwrite(fn, g3im)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
"""
|
||||
Now only feasible for trackA_XX
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root_dir', default="../datasets/icdar2019/at_trackA_archival")
|
||||
args = parser.parse_args()
|
||||
|
||||
for fdname in os.listdir(args.root_dir):
|
||||
if fdname.endswith(".json"):
|
||||
continue
|
||||
ffdname = os.path.join(args.root_dir, fdname)
|
||||
for file in tqdm.tqdm(os.listdir(ffdname)):
|
||||
if file.endswith(".xml"):
|
||||
continue
|
||||
ffile = os.path.join(ffdname, file)
|
||||
convert(ffile)
|
||||
@@ -0,0 +1,104 @@
|
||||
import os
|
||||
from PIL import Image
|
||||
import xml.etree.ElementTree as ET
|
||||
import numpy as np
|
||||
import json
|
||||
from PIL import Image
|
||||
from shutil import copyfile
|
||||
|
||||
|
||||
def convert(ROOT, TRACK, SPLIT):
|
||||
coco_data = {
|
||||
"images": [],
|
||||
"annotations": [],
|
||||
"categories": [{"id": 1, "name": "table"}, ],
|
||||
}
|
||||
DATA_DIR = f"{ROOT}/{TRACK}/{SPLIT}"
|
||||
prefix = "cTDaR_t0" if TRACK == "trackA_archival" else "cTDaR_t1"
|
||||
print(TRACK, SPLIT, prefix)
|
||||
table_count = 0
|
||||
for file in sorted(os.listdir(DATA_DIR)):
|
||||
if file.startswith(prefix) and file.endswith(".jpg"):
|
||||
img = Image.open(os.path.join(DATA_DIR, file))
|
||||
coco_data["images"].append(
|
||||
{
|
||||
"file_name": file,
|
||||
"height": img.height,
|
||||
"width": img.width,
|
||||
"id": int(file[7:-4]),
|
||||
}
|
||||
)
|
||||
elif file.startswith(prefix) and file.endswith(".xml"):
|
||||
# print(file)
|
||||
tree = ET.parse(os.path.join(DATA_DIR, file))
|
||||
root = tree.getroot()
|
||||
assert len(root.findall("./table/Coords")) > 0
|
||||
for table_id in range(len(root.findall("./table/Coords"))):
|
||||
four_points = root.findall("./table/Coords")[table_id].attrib["points"]
|
||||
four_points = list(map(lambda x: x.split(","), four_points.split()))
|
||||
four_points = [[int(j) for j in i] for i in four_points]
|
||||
segmentation = [j for i in four_points for j in i]
|
||||
bbox = [
|
||||
four_points[0][0],
|
||||
four_points[0][1],
|
||||
four_points[2][0] - four_points[0][0],
|
||||
four_points[2][1] - four_points[0][1],
|
||||
]
|
||||
coco_data["annotations"].append(
|
||||
{
|
||||
"segmentation": [segmentation],
|
||||
"area": bbox[2] * bbox[3],
|
||||
"iscrowd": 0,
|
||||
"image_id": int(file[7:-4]),
|
||||
"bbox": bbox,
|
||||
"category_id": 1,
|
||||
"id": table_count,
|
||||
}
|
||||
)
|
||||
table_count += 1
|
||||
|
||||
with open(f"{ROOT}/{TRACK}/{SPLIT}.json", "w") as f:
|
||||
json.dump(coco_data, f)
|
||||
|
||||
|
||||
def clean_img(DATA_DIR):
|
||||
for file in sorted(os.listdir(DATA_DIR)):
|
||||
if file.endswith(".JPG"):
|
||||
os.rename(os.path.join(DATA_DIR, file), os.path.join(DATA_DIR, file.replace(".JPG", ".jpg")))
|
||||
elif file.endswith(".TIFF"):
|
||||
img = Image.open(os.path.join(DATA_DIR, file))
|
||||
img.save(os.path.join(DATA_DIR, file.replace(".TIFF", ".jpg")))
|
||||
os.remove(os.path.join(DATA_DIR, file))
|
||||
elif file.endswith(".png"):
|
||||
img = Image.open(os.path.join(DATA_DIR, file))
|
||||
img.save(os.path.join(DATA_DIR, file.replace(".png", ".jpg")))
|
||||
os.remove(os.path.join(DATA_DIR, file))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root_dir', required=True)
|
||||
parser.add_argument('--target_dir', required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
test_data_dir = os.path.join(args.root_dir, 'test', 'TRACKA')
|
||||
test_gt_dir = os.path.join(args.root_dir, 'test_ground_truth', 'TRACKA')
|
||||
training_data_dir = os.path.join(args.root_dir, 'training', 'TRACKA', 'ground_truth')
|
||||
|
||||
raw_datas = {"train": [training_data_dir], "test": [test_data_dir, test_gt_dir]}
|
||||
|
||||
TRACKS = ["trackA_modern", "trackA_archival"]
|
||||
SPLITS = ["train", "test"]
|
||||
for track in TRACKS:
|
||||
prefix = "cTDaR_t0" if track == "trackA_archival" else "cTDaR_t1"
|
||||
for split in SPLITS:
|
||||
os.makedirs(os.path.join(args.target_dir, track, split))
|
||||
for source_dir in raw_datas[split]:
|
||||
for fn in os.listdir(source_dir):
|
||||
if fn.startswith(prefix):
|
||||
ffn = os.path.join(source_dir, fn)
|
||||
copyfile(ffn, os.path.join(args.target_dir, track, split, fn))
|
||||
clean_img(os.path.join(args.target_dir, track, split))
|
||||
convert(args.target_dir, track, split)
|
||||
@@ -0,0 +1,16 @@
|
||||
# --------------------------------------------------------------------------------
|
||||
# MPViT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
from .config import add_vit_config
|
||||
from .backbone import build_vit_fpn_backbone
|
||||
from .dataset_mapper import DetrDatasetMapper
|
||||
from .mycheckpointer import MyDetectionCheckpointer
|
||||
from .icdar_evaluation import ICDAREvaluator
|
||||
from .mytrainer import MyTrainer
|
||||
from .table_evaluation import calc_table_score
|
||||
@@ -0,0 +1,156 @@
|
||||
# --------------------------------------------------------------------------------
|
||||
# VIT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# --------------------------------------------------------------------------------
|
||||
# References:
|
||||
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# CoaT: https://github.com/mlpc-ucsd/CoaT
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
from detectron2.layers import (
|
||||
ShapeSpec,
|
||||
)
|
||||
from detectron2.modeling import Backbone, BACKBONE_REGISTRY, FPN
|
||||
from detectron2.modeling.backbone.fpn import LastLevelP6P7, LastLevelMaxPool
|
||||
|
||||
from .beit import beit_base_patch16, dit_base_patch16, dit_large_patch16, beit_large_patch16
|
||||
from .deit import deit_base_patch16, mae_base_patch16
|
||||
|
||||
__all__ = [
|
||||
"build_vit_fpn_backbone",
|
||||
]
|
||||
|
||||
|
||||
class VIT_Backbone(Backbone):
|
||||
"""
|
||||
Implement VIT backbone.
|
||||
"""
|
||||
|
||||
def __init__(self, name, out_features, drop_path, img_size, pos_type, model_kwargs):
|
||||
super().__init__()
|
||||
self._out_features = out_features
|
||||
if 'base' in name:
|
||||
self._out_feature_strides = {"layer3": 4, "layer5": 8, "layer7": 16, "layer11": 32}
|
||||
else:
|
||||
self._out_feature_strides = {"layer7": 4, "layer11": 8, "layer15": 16, "layer23": 32}
|
||||
|
||||
if name == 'beit_base_patch16':
|
||||
model_func = beit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == 'dit_base_patch16':
|
||||
model_func = dit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "deit_base_patch16":
|
||||
model_func = deit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "mae_base_patch16":
|
||||
model_func = mae_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "dit_large_patch16":
|
||||
model_func = dit_large_patch16
|
||||
self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
|
||||
elif name == "beit_large_patch16":
|
||||
model_func = beit_large_patch16
|
||||
self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
|
||||
else:
|
||||
raise ValueError("Unsupported VIT name yet.")
|
||||
|
||||
if 'beit' in name or 'dit' in name:
|
||||
if pos_type == "abs":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_abs_pos_emb=True,
|
||||
**model_kwargs)
|
||||
elif pos_type == "shared_rel":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_shared_rel_pos_bias=True,
|
||||
**model_kwargs)
|
||||
elif pos_type == "rel":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_rel_pos_bias=True,
|
||||
**model_kwargs)
|
||||
else:
|
||||
raise ValueError()
|
||||
else:
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
**model_kwargs)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
|
||||
|
||||
Returns:
|
||||
dict[str->Tensor]: names and the corresponding features
|
||||
"""
|
||||
assert x.dim() == 4, f"VIT takes an input of shape (N, C, H, W). Got {x.shape} instead!"
|
||||
return self.backbone.forward_features(x)
|
||||
|
||||
def output_shape(self):
|
||||
return {
|
||||
name: ShapeSpec(
|
||||
channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
|
||||
)
|
||||
for name in self._out_features
|
||||
}
|
||||
|
||||
|
||||
def build_VIT_backbone(cfg):
|
||||
"""
|
||||
Create a VIT instance from config.
|
||||
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
A VIT backbone instance.
|
||||
"""
|
||||
# fmt: off
|
||||
name = cfg.MODEL.VIT.NAME
|
||||
out_features = cfg.MODEL.VIT.OUT_FEATURES
|
||||
drop_path = cfg.MODEL.VIT.DROP_PATH
|
||||
img_size = cfg.MODEL.VIT.IMG_SIZE
|
||||
pos_type = cfg.MODEL.VIT.POS_TYPE
|
||||
|
||||
model_kwargs = eval(str(cfg.MODEL.VIT.MODEL_KWARGS).replace("`", ""))
|
||||
|
||||
return VIT_Backbone(name, out_features, drop_path, img_size, pos_type, model_kwargs)
|
||||
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def build_vit_fpn_backbone(cfg, input_shape: ShapeSpec):
|
||||
"""
|
||||
Create a VIT w/ FPN backbone.
|
||||
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_VIT_backbone(cfg)
|
||||
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm=cfg.MODEL.FPN.NORM,
|
||||
top_block=LastLevelMaxPool(),
|
||||
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||
)
|
||||
return backbone
|
||||
@@ -0,0 +1,671 @@
|
||||
""" 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 warnings
|
||||
import math
|
||||
import torch
|
||||
from functools import partial
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint as checkpoint
|
||||
from timm.models.layers import drop_path, to_2tuple, trunc_normal_
|
||||
|
||||
|
||||
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 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.):
|
||||
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)
|
||||
# commit this for the orignal BERT implement
|
||||
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., window_size=None, attn_head_dim=None):
|
||||
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
|
||||
# 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, 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
|
||||
|
||||
if window_size:
|
||||
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)
|
||||
|
||||
# trunc_normal_(self.relative_position_bias_table, std=.0)
|
||||
else:
|
||||
self.window_size = None
|
||||
self.relative_position_bias_table = None
|
||||
self.relative_position_index = None
|
||||
|
||||
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, training_window_size=None):
|
||||
B, N, C = x.shape
|
||||
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 = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
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)
|
||||
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 self.relative_position_bias_table is not None:
|
||||
if training_window_size == self.window_size:
|
||||
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)
|
||||
else:
|
||||
training_window_size = tuple(training_window_size.tolist())
|
||||
new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
|
||||
# new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
|
||||
new_relative_position_bias_table = F.interpolate(
|
||||
self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
|
||||
2 * self.window_size[0] - 1,
|
||||
2 * self.window_size[1] - 1),
|
||||
size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
|
||||
align_corners=False)
|
||||
new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
|
||||
new_num_relative_distance - 3).permute(
|
||||
1, 0)
|
||||
new_relative_position_bias_table = torch.cat(
|
||||
[new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(training_window_size[0])
|
||||
coords_w = torch.arange(training_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] += training_window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += training_window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(training_window_size[0] * training_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:] = new_num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = new_num_relative_distance - 2
|
||||
relative_position_index[0, 0] = new_num_relative_distance - 1
|
||||
|
||||
relative_position_bias = \
|
||||
new_relative_position_bias_table[relative_position_index.view(-1)].view(
|
||||
training_window_size[0] * training_window_size[1] + 1,
|
||||
training_window_size[0] * training_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)
|
||||
|
||||
if rel_pos_bias is not None:
|
||||
attn = attn + rel_pos_bias
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
|
||||
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., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
|
||||
window_size=None, attn_head_dim=None):
|
||||
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, window_size=window_size, attn_head_dim=attn_head_dim)
|
||||
# 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)
|
||||
|
||||
if init_values is not None:
|
||||
self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
|
||||
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, rel_pos_bias=None, training_window_size=None):
|
||||
if self.gamma_1 is None:
|
||||
x = x + self.drop_path(
|
||||
self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, training_window_size=training_window_size))
|
||||
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
||||
else:
|
||||
x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias,
|
||||
training_window_size=training_window_size))
|
||||
x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
""" Image to Patch Embedding
|
||||
"""
|
||||
|
||||
def __init__(self, img_size=[224, 224], patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
|
||||
self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
self.num_patches_w = self.patch_shape[0]
|
||||
self.num_patches_h = self.patch_shape[1]
|
||||
# the so-called patch_shape is the patch shape during pre-training
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.num_patches = num_patches
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x, position_embedding=None, **kwargs):
|
||||
# FIXME look at relaxing size constraints
|
||||
# assert H == self.img_size[0] and W == self.img_size[1], \
|
||||
# f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
||||
x = self.proj(x)
|
||||
Hp, Wp = x.shape[2], x.shape[3]
|
||||
|
||||
if position_embedding is not None:
|
||||
# interpolate the position embedding to the corresponding size
|
||||
position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3,
|
||||
1, 2)
|
||||
position_embedding = F.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
|
||||
x = x + position_embedding
|
||||
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
return x, (Hp, Wp)
|
||||
|
||||
|
||||
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, 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 RelativePositionBias(nn.Module):
|
||||
|
||||
def __init__(self, window_size, num_heads):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.num_heads = num_heads
|
||||
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)
|
||||
|
||||
# trunc_normal_(self.relative_position_bias_table, std=.02)
|
||||
|
||||
def forward(self, training_window_size):
|
||||
if training_window_size == self.window_size:
|
||||
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
|
||||
else:
|
||||
training_window_size = tuple(training_window_size.tolist())
|
||||
new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
|
||||
# new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
|
||||
new_relative_position_bias_table = F.interpolate(
|
||||
self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
|
||||
2 * self.window_size[0] - 1,
|
||||
2 * self.window_size[1] - 1),
|
||||
size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
|
||||
align_corners=False)
|
||||
new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
|
||||
new_num_relative_distance - 3).permute(
|
||||
1, 0)
|
||||
new_relative_position_bias_table = torch.cat(
|
||||
[new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(training_window_size[0])
|
||||
coords_w = torch.arange(training_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] += training_window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += training_window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(training_window_size[0] * training_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:] = new_num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = new_num_relative_distance - 2
|
||||
relative_position_index[0, 0] = new_num_relative_distance - 1
|
||||
|
||||
relative_position_bias = \
|
||||
new_relative_position_bias_table[relative_position_index.view(-1)].view(
|
||||
training_window_size[0] * training_window_size[1] + 1,
|
||||
training_window_size[0] * training_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
|
||||
|
||||
return relative_position_bias
|
||||
|
||||
|
||||
class BEiT(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
img_size=[224, 224],
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
num_classes=80,
|
||||
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=None,
|
||||
init_values=None,
|
||||
use_abs_pos_emb=False,
|
||||
use_rel_pos_bias=False,
|
||||
use_shared_rel_pos_bias=False,
|
||||
use_checkpoint=True,
|
||||
pretrained=None,
|
||||
out_features=None,
|
||||
):
|
||||
|
||||
super(BEiT, self).__init__()
|
||||
|
||||
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
|
||||
self.num_classes = num_classes
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
if hybrid_backbone is not None:
|
||||
self.patch_embed = HybridEmbed(
|
||||
hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
else:
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
self.out_features = out_features
|
||||
self.out_indices = [int(name[5:]) for name in out_features]
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
# self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
if use_abs_pos_emb:
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
|
||||
else:
|
||||
self.pos_embed = None
|
||||
self.pos_drop = nn.Dropout(p=drop_rate)
|
||||
|
||||
self.use_shared_rel_pos_bias = use_shared_rel_pos_bias
|
||||
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
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
self.use_rel_pos_bias = use_rel_pos_bias
|
||||
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,
|
||||
init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)
|
||||
for i in range(depth)])
|
||||
|
||||
# trunc_normal_(self.mask_token, std=.02)
|
||||
|
||||
if patch_size == 16:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
# nn.SyncBatchNorm(embed_dim),
|
||||
nn.BatchNorm2d(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn3 = nn.Identity()
|
||||
|
||||
self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
elif patch_size == 8:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Identity()
|
||||
|
||||
self.fpn3 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn4 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=4, stride=4),
|
||||
)
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
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 init_weights(self):
|
||||
"""Initialize the weights in backbone.
|
||||
|
||||
Args:
|
||||
pretrained (str, optional): Path to pre-trained weights.
|
||||
Defaults to None.
|
||||
"""
|
||||
logger = get_root_logger()
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
if self.init_cfg is None:
|
||||
logger.warn(f'No pre-trained weights for '
|
||||
f'{self.__class__.__name__}, '
|
||||
f'training start from scratch')
|
||||
else:
|
||||
assert 'checkpoint' in self.init_cfg, f'Only support ' \
|
||||
f'specify `Pretrained` in ' \
|
||||
f'`init_cfg` in ' \
|
||||
f'{self.__class__.__name__} '
|
||||
logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
|
||||
load_checkpoint(self,
|
||||
filename=self.init_cfg['checkpoint'],
|
||||
strict=False,
|
||||
logger=logger,
|
||||
beit_spec_expand_rel_pos = self.use_rel_pos_bias,
|
||||
)
|
||||
'''
|
||||
|
||||
def get_num_layers(self):
|
||||
return len(self.blocks)
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'pos_embed', 'cls_token'}
|
||||
|
||||
def forward_features(self, x):
|
||||
B, C, H, W = x.shape
|
||||
x, (Hp, Wp) = self.patch_embed(x, self.pos_embed[:, 1:, :] if self.pos_embed is not None else None)
|
||||
# Hp, Wp are HW for patches
|
||||
batch_size, seq_len, _ = x.size()
|
||||
|
||||
cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
if self.pos_embed is not None:
|
||||
cls_tokens = cls_tokens + self.pos_embed[:, :1, :]
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
x = self.pos_drop(x)
|
||||
|
||||
features = []
|
||||
training_window_size = torch.tensor([Hp, Wp])
|
||||
|
||||
rel_pos_bias = self.rel_pos_bias(training_window_size) if self.rel_pos_bias is not None else None
|
||||
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if self.use_checkpoint:
|
||||
x = checkpoint.checkpoint(blk, x, rel_pos_bias, training_window_size)
|
||||
else:
|
||||
x = blk(x, rel_pos_bias=rel_pos_bias, training_window_size=training_window_size)
|
||||
if i in self.out_indices:
|
||||
xp = x[:, 1:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
|
||||
features.append(xp.contiguous())
|
||||
|
||||
ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
|
||||
for i in range(len(features)):
|
||||
features[i] = ops[i](features[i])
|
||||
|
||||
feat_out = {}
|
||||
|
||||
for name, value in zip(self.out_features, features):
|
||||
feat_out[name] = value
|
||||
|
||||
return feat_out
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
return x
|
||||
|
||||
|
||||
def beit_base_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=None,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def beit_large_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=None,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def dit_base_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=0.1,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def dit_large_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=1e-5,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = BEiT(use_checkpoint=True, use_shared_rel_pos_bias=True)
|
||||
model = model.to("cuda:0")
|
||||
input1 = torch.rand(2, 3, 512, 762).to("cuda:0")
|
||||
input2 = torch.rand(2, 3, 800, 1200).to("cuda:0")
|
||||
input3 = torch.rand(2, 3, 720, 1000).to("cuda:0")
|
||||
output1 = model(input1)
|
||||
output2 = model(input2)
|
||||
output3 = model(input3)
|
||||
print("all done")
|
||||
@@ -0,0 +1,32 @@
|
||||
from detectron2.config import CfgNode as CN
|
||||
|
||||
|
||||
def add_vit_config(cfg):
|
||||
"""
|
||||
Add config for VIT.
|
||||
"""
|
||||
_C = cfg
|
||||
|
||||
_C.MODEL.VIT = CN()
|
||||
|
||||
# CoaT model name.
|
||||
_C.MODEL.VIT.NAME = ""
|
||||
|
||||
# Output features from CoaT backbone.
|
||||
_C.MODEL.VIT.OUT_FEATURES = ["layer3", "layer5", "layer7", "layer11"]
|
||||
|
||||
_C.MODEL.VIT.IMG_SIZE = [224, 224]
|
||||
|
||||
_C.MODEL.VIT.POS_TYPE = "shared_rel"
|
||||
|
||||
_C.MODEL.VIT.DROP_PATH = 0.
|
||||
|
||||
_C.MODEL.VIT.MODEL_KWARGS = "{}"
|
||||
|
||||
_C.SOLVER.OPTIMIZER = "ADAMW"
|
||||
|
||||
_C.SOLVER.BACKBONE_MULTIPLIER = 1.0
|
||||
|
||||
_C.AUG = CN()
|
||||
|
||||
_C.AUG.DETR = False
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
# from https://github.com/facebookresearch/detr/blob/main/d2/detr/dataset_mapper.py
|
||||
|
||||
|
||||
import copy
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from detectron2.data import detection_utils as utils
|
||||
from detectron2.data import transforms as T
|
||||
|
||||
__all__ = ["DetrDatasetMapper"]
|
||||
|
||||
|
||||
def build_transform_gen(cfg, is_train):
|
||||
"""
|
||||
Create a list of :class:`TransformGen` from config.
|
||||
Returns:
|
||||
list[TransformGen]
|
||||
"""
|
||||
if is_train:
|
||||
min_size = cfg.INPUT.MIN_SIZE_TRAIN
|
||||
max_size = cfg.INPUT.MAX_SIZE_TRAIN
|
||||
sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING
|
||||
else:
|
||||
min_size = cfg.INPUT.MIN_SIZE_TEST
|
||||
max_size = cfg.INPUT.MAX_SIZE_TEST
|
||||
sample_style = "choice"
|
||||
if sample_style == "range":
|
||||
assert len(min_size) == 2, "more than 2 ({}) min_size(s) are provided for ranges".format(len(min_size))
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
tfm_gens = []
|
||||
if is_train:
|
||||
tfm_gens.append(T.RandomFlip())
|
||||
tfm_gens.append(T.ResizeShortestEdge(min_size, max_size, sample_style))
|
||||
if is_train:
|
||||
logger.info("TransformGens used in training: " + str(tfm_gens))
|
||||
return tfm_gens
|
||||
|
||||
|
||||
class DetrDatasetMapper:
|
||||
"""
|
||||
A callable which takes a dataset dict in Detectron2 Dataset format,
|
||||
and map it into a format used by DETR.
|
||||
|
||||
The callable currently does the following:
|
||||
|
||||
1. Read the image from "file_name"
|
||||
2. Applies geometric transforms to the image and annotation
|
||||
3. Find and applies suitable cropping to the image and annotation
|
||||
4. Prepare image and annotation to Tensors
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, is_train=True):
|
||||
if cfg.INPUT.CROP.ENABLED and is_train:
|
||||
self.crop_gen = [
|
||||
T.ResizeShortestEdge([400, 500, 600], sample_style="choice"),
|
||||
T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE),
|
||||
]
|
||||
else:
|
||||
self.crop_gen = None
|
||||
|
||||
self.mask_on = cfg.MODEL.MASK_ON
|
||||
self.tfm_gens = build_transform_gen(cfg, is_train)
|
||||
logging.getLogger(__name__).info(
|
||||
"Full TransformGens used in training: {}, crop: {}".format(str(self.tfm_gens), str(self.crop_gen))
|
||||
)
|
||||
|
||||
self.img_format = cfg.INPUT.FORMAT
|
||||
self.is_train = is_train
|
||||
|
||||
def __call__(self, dataset_dict):
|
||||
"""
|
||||
Args:
|
||||
dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
|
||||
|
||||
Returns:
|
||||
dict: a format that builtin models in detectron2 accept
|
||||
"""
|
||||
dataset_dict = copy.deepcopy(dataset_dict) # it will be modified by code below
|
||||
image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
|
||||
utils.check_image_size(dataset_dict, image)
|
||||
|
||||
if self.crop_gen is None:
|
||||
image, transforms = T.apply_transform_gens(self.tfm_gens, image)
|
||||
else:
|
||||
if np.random.rand() > 0.5:
|
||||
image, transforms = T.apply_transform_gens(self.tfm_gens, image)
|
||||
else:
|
||||
image, transforms = T.apply_transform_gens(
|
||||
self.tfm_gens[:-1] + self.crop_gen + self.tfm_gens[-1:], image
|
||||
)
|
||||
|
||||
image_shape = image.shape[:2] # h, w
|
||||
|
||||
# Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,
|
||||
# but not efficient on large generic data structures due to the use of pickle & mp.Queue.
|
||||
# Therefore it's important to use torch.Tensor.
|
||||
dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
|
||||
|
||||
if not self.is_train:
|
||||
# USER: Modify this if you want to keep them for some reason.
|
||||
dataset_dict.pop("annotations", None)
|
||||
return dataset_dict
|
||||
|
||||
if "annotations" in dataset_dict:
|
||||
# USER: Modify this if you want to keep them for some reason.
|
||||
for anno in dataset_dict["annotations"]:
|
||||
if not self.mask_on:
|
||||
anno.pop("segmentation", None)
|
||||
anno.pop("keypoints", None)
|
||||
|
||||
# USER: Implement additional transformations if you have other types of data
|
||||
annos = [
|
||||
utils.transform_instance_annotations(obj, transforms, image_shape)
|
||||
for obj in dataset_dict.pop("annotations")
|
||||
if obj.get("iscrowd", 0) == 0
|
||||
]
|
||||
instances = utils.annotations_to_instances(annos, image_shape)
|
||||
dataset_dict["instances"] = utils.filter_empty_instances(instances)
|
||||
return dataset_dict
|
||||
@@ -0,0 +1,476 @@
|
||||
"""
|
||||
Mostly copy-paste from DINO and timm library:
|
||||
https://github.com/facebookresearch/dino
|
||||
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
|
||||
"""
|
||||
import warnings
|
||||
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint as checkpoint
|
||||
from timm.models.layers import trunc_normal_, drop_path, to_2tuple
|
||||
from functools import partial
|
||||
|
||||
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 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.):
|
||||
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
|
||||
q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads,
|
||||
C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
|
||||
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, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
|
||||
self.window_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
|
||||
self.num_patches_w, self.num_patches_h = self.window_size
|
||||
|
||||
self.num_patches = self.window_size[0] * self.window_size[1]
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim,
|
||||
kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.proj(x)
|
||||
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 ViT(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
model_name='vit_base_patch16_224',
|
||||
img_size=384,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
num_classes=19,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
qk_scale=None,
|
||||
drop_rate=0.1,
|
||||
attn_drop_rate=0.,
|
||||
drop_path_rate=0.,
|
||||
hybrid_backbone=None,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
norm_cfg=None,
|
||||
pos_embed_interp=False,
|
||||
random_init=False,
|
||||
align_corners=False,
|
||||
use_checkpoint=False,
|
||||
num_extra_tokens=1,
|
||||
out_features=None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
super(ViT, self).__init__()
|
||||
self.model_name = model_name
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.in_chans = in_chans
|
||||
self.embed_dim = embed_dim
|
||||
self.depth = depth
|
||||
self.num_heads = num_heads
|
||||
self.num_classes = num_classes
|
||||
self.mlp_ratio = mlp_ratio
|
||||
self.qkv_bias = qkv_bias
|
||||
self.qk_scale = qk_scale
|
||||
self.drop_rate = drop_rate
|
||||
self.attn_drop_rate = attn_drop_rate
|
||||
self.drop_path_rate = drop_path_rate
|
||||
self.hybrid_backbone = hybrid_backbone
|
||||
self.norm_layer = norm_layer
|
||||
self.norm_cfg = norm_cfg
|
||||
self.pos_embed_interp = pos_embed_interp
|
||||
self.random_init = random_init
|
||||
self.align_corners = align_corners
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.num_extra_tokens = num_extra_tokens
|
||||
self.out_features = out_features
|
||||
self.out_indices = [int(name[5:]) for name in out_features]
|
||||
|
||||
# self.num_stages = self.depth
|
||||
# self.out_indices = tuple(range(self.num_stages))
|
||||
|
||||
if self.hybrid_backbone is not None:
|
||||
self.patch_embed = HybridEmbed(
|
||||
self.hybrid_backbone, img_size=self.img_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
|
||||
else:
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=self.img_size, patch_size=self.patch_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
|
||||
self.num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
|
||||
|
||||
if self.num_extra_tokens == 2:
|
||||
self.dist_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
|
||||
|
||||
self.pos_embed = nn.Parameter(torch.zeros(
|
||||
1, self.num_patches + self.num_extra_tokens, self.embed_dim))
|
||||
self.pos_drop = nn.Dropout(p=self.drop_rate)
|
||||
|
||||
# self.num_extra_tokens = self.pos_embed.shape[-2] - self.num_patches
|
||||
dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate,
|
||||
self.depth)] # stochastic depth decay rule
|
||||
self.blocks = nn.ModuleList([
|
||||
Block(
|
||||
dim=self.embed_dim, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, qkv_bias=self.qkv_bias,
|
||||
qk_scale=self.qk_scale,
|
||||
drop=self.drop_rate, attn_drop=self.attn_drop_rate, drop_path=dpr[i], norm_layer=self.norm_layer)
|
||||
for i in range(self.depth)])
|
||||
|
||||
# NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
|
||||
# self.repr = nn.Linear(embed_dim, representation_size)
|
||||
# self.repr_act = nn.Tanh()
|
||||
|
||||
if patch_size == 16:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
nn.SyncBatchNorm(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn3 = nn.Identity()
|
||||
|
||||
self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
elif patch_size == 8:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Identity()
|
||||
|
||||
self.fpn3 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn4 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=4, stride=4),
|
||||
)
|
||||
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
if self.num_extra_tokens==2:
|
||||
trunc_normal_(self.dist_token, std=0.2)
|
||||
self.apply(self._init_weights)
|
||||
# self.fix_init_weight()
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
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 init_weights(self):
|
||||
logger = get_root_logger()
|
||||
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
if self.init_cfg is None:
|
||||
logger.warn(f'No pre-trained weights for '
|
||||
f'{self.__class__.__name__}, '
|
||||
f'training start from scratch')
|
||||
else:
|
||||
assert 'checkpoint' in self.init_cfg, f'Only support ' \
|
||||
f'specify `Pretrained` in ' \
|
||||
f'`init_cfg` in ' \
|
||||
f'{self.__class__.__name__} '
|
||||
logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
|
||||
load_checkpoint(self, filename=self.init_cfg['checkpoint'], strict=False, logger=logger)
|
||||
'''
|
||||
|
||||
def get_num_layers(self):
|
||||
return len(self.blocks)
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'pos_embed', 'cls_token'}
|
||||
|
||||
def _conv_filter(self, 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
|
||||
|
||||
def to_2D(self, x):
|
||||
n, hw, c = x.shape
|
||||
h = w = int(math.sqrt(hw))
|
||||
x = x.transpose(1, 2).reshape(n, c, h, w)
|
||||
return x
|
||||
|
||||
def to_1D(self, x):
|
||||
n, c, h, w = x.shape
|
||||
x = x.reshape(n, c, -1).transpose(1, 2)
|
||||
return x
|
||||
|
||||
def interpolate_pos_encoding(self, x, w, h):
|
||||
npatch = x.shape[1] - self.num_extra_tokens
|
||||
N = self.pos_embed.shape[1] - self.num_extra_tokens
|
||||
if npatch == N and w == h:
|
||||
return self.pos_embed
|
||||
|
||||
class_ORdist_pos_embed = self.pos_embed[:, 0:self.num_extra_tokens]
|
||||
|
||||
patch_pos_embed = self.pos_embed[:, self.num_extra_tokens:]
|
||||
|
||||
dim = x.shape[-1]
|
||||
w0 = w // self.patch_embed.patch_size[0]
|
||||
h0 = h // self.patch_embed.patch_size[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 = w0 + 0.1, h0 + 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=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
|
||||
mode='bicubic',
|
||||
)
|
||||
assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
|
||||
return torch.cat((class_ORdist_pos_embed, patch_pos_embed), dim=1)
|
||||
|
||||
def prepare_tokens(self, x, mask=None):
|
||||
B, nc, w, h = x.shape
|
||||
# patch linear embedding
|
||||
x = self.patch_embed(x)
|
||||
|
||||
# mask image modeling
|
||||
if mask is not None:
|
||||
x = self.mask_model(x, mask)
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
|
||||
# add the [CLS] token to the embed patch tokens
|
||||
all_tokens = [self.cls_token.expand(B, -1, -1)]
|
||||
|
||||
if self.num_extra_tokens == 2:
|
||||
dist_tokens = self.dist_token.expand(B, -1, -1)
|
||||
all_tokens.append(dist_tokens)
|
||||
all_tokens.append(x)
|
||||
|
||||
x = torch.cat(all_tokens, dim=1)
|
||||
|
||||
# add positional encoding to each token
|
||||
x = x + self.interpolate_pos_encoding(x, w, h)
|
||||
|
||||
return self.pos_drop(x)
|
||||
|
||||
def forward_features(self, x):
|
||||
# print(f"==========shape of x is {x.shape}==========")
|
||||
B, _, H, W = x.shape
|
||||
Hp, Wp = H // self.patch_size, W // self.patch_size
|
||||
x = self.prepare_tokens(x)
|
||||
|
||||
features = []
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if self.use_checkpoint:
|
||||
x = checkpoint.checkpoint(blk, x)
|
||||
else:
|
||||
x = blk(x)
|
||||
if i in self.out_indices:
|
||||
xp = x[:, self.num_extra_tokens:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
|
||||
features.append(xp.contiguous())
|
||||
|
||||
ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
|
||||
for i in range(len(features)):
|
||||
features[i] = ops[i](features[i])
|
||||
|
||||
feat_out = {}
|
||||
|
||||
for name, value in zip(self.out_features, features):
|
||||
feat_out[name] = value
|
||||
|
||||
return feat_out
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
return x
|
||||
|
||||
|
||||
def deit_base_patch16(pretrained=False, **kwargs):
|
||||
model = ViT(
|
||||
patch_size=16,
|
||||
drop_rate=0.,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
num_classes=1000,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
use_checkpoint=True,
|
||||
num_extra_tokens=2,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def mae_base_patch16(pretrained=False, **kwargs):
|
||||
model = ViT(
|
||||
patch_size=16,
|
||||
drop_rate=0.,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
num_classes=1000,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
use_checkpoint=True,
|
||||
num_extra_tokens=1,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
@@ -0,0 +1,100 @@
|
||||
import copy
|
||||
import itertools
|
||||
import os
|
||||
import os.path as osp
|
||||
import shutil
|
||||
from collections import OrderedDict
|
||||
from xml.dom.minidom import Document
|
||||
|
||||
import detectron2.utils.comm as comm
|
||||
import torch
|
||||
from detectron2.evaluation import COCOEvaluator
|
||||
from detectron2.utils.file_io import PathManager
|
||||
|
||||
from .table_evaluation.evaluate import calc_table_score
|
||||
|
||||
|
||||
class ICDAREvaluator(COCOEvaluator):
|
||||
def evaluate(self, img_ids=None):
|
||||
"""
|
||||
Args:
|
||||
img_ids: a list of image IDs to evaluate on. Default to None for the whole dataset
|
||||
"""
|
||||
if self._distributed:
|
||||
comm.synchronize()
|
||||
predictions = comm.gather(self._predictions, dst=0)
|
||||
predictions = list(itertools.chain(*predictions))
|
||||
|
||||
if not comm.is_main_process():
|
||||
return {}
|
||||
else:
|
||||
predictions = self._predictions
|
||||
|
||||
if len(predictions) == 0:
|
||||
self._logger.warning("[COCOEvaluator] Did not receive valid predictions.")
|
||||
return {}
|
||||
|
||||
if self._output_dir:
|
||||
PathManager.mkdirs(self._output_dir)
|
||||
file_path = os.path.join(self._output_dir, "instances_predictions.pth")
|
||||
with PathManager.open(file_path, "wb") as f:
|
||||
torch.save(predictions, f)
|
||||
|
||||
self._results = OrderedDict()
|
||||
if "proposals" in predictions[0]:
|
||||
self._eval_box_proposals(predictions)
|
||||
if "instances" in predictions[0]:
|
||||
self._eval_predictions(predictions, img_ids=img_ids)
|
||||
self.evaluate_table(predictions)
|
||||
# Copy so the caller can do whatever with results
|
||||
return copy.deepcopy(self._results)
|
||||
|
||||
def evaluate_table(self, predictions):
|
||||
xml_dir = self.convert_to_xml(predictions)
|
||||
results = calc_table_score(xml_dir)
|
||||
self._results["wF1"] = results['wF1']
|
||||
|
||||
def convert_to_xml(self, predictions):
|
||||
output_dir = osp.join(self._output_dir, "xml_results")
|
||||
if os.path.exists(output_dir):
|
||||
shutil.rmtree(output_dir)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
|
||||
results_dict = {}
|
||||
for result in coco_results:
|
||||
if result["score"] < 0.7:
|
||||
continue
|
||||
image_id = result["image_id"]
|
||||
if image_id not in results_dict:
|
||||
results_dict[image_id] = []
|
||||
|
||||
results_dict[image_id].append(result)
|
||||
|
||||
for image_id, tables in results_dict.items():
|
||||
file_name = f"cTDaR_t{image_id:05d}.jpg"
|
||||
doc = Document()
|
||||
root = doc.createElement('document')
|
||||
root.setAttribute('filename', file_name)
|
||||
doc.appendChild(root)
|
||||
for table_id, table in enumerate(tables, start=1):
|
||||
nodeManager = doc.createElement('table')
|
||||
nodeManager.setAttribute('id', str(table_id))
|
||||
bbox = list(map(int, table['bbox']))
|
||||
bbox_str = '{},{} {},{} {},{} {},{}'.format(bbox[0], bbox[1],
|
||||
bbox[0], bbox[1] + bbox[3],
|
||||
bbox[0] + bbox[2], bbox[1] + bbox[3],
|
||||
bbox[0] + bbox[2], bbox[1])
|
||||
nodeCoords = doc.createElement('Coords')
|
||||
nodeCoords.setAttribute('points', bbox_str)
|
||||
nodeManager.appendChild(nodeCoords)
|
||||
root.appendChild(nodeManager)
|
||||
filename = '{}-result.xml'.format(file_name[:-4])
|
||||
fp = open(os.path.join(output_dir, filename), 'w')
|
||||
doc.writexml(fp, indent='', addindent='\t', newl='\n', encoding="utf-8")
|
||||
fp.flush()
|
||||
fp.close()
|
||||
return output_dir
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pass
|
||||
@@ -0,0 +1,257 @@
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
|
||||
from typing import Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from fvcore.common.checkpoint import _IncompatibleKeys, _strip_prefix_if_present, TORCH_VERSION, quantization, \
|
||||
ObserverBase, FakeQuantizeBase
|
||||
from torch import distributed as dist
|
||||
from scipy import interpolate
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def append_prefix(k):
|
||||
prefix = 'backbone.bottom_up.backbone.'
|
||||
return prefix + k if not k.startswith(prefix) else k
|
||||
|
||||
|
||||
def modify_ckpt_state(model, state_dict, logger=None):
|
||||
# reshape absolute position embedding for Swin
|
||||
if state_dict.get(append_prefix('absolute_pos_embed')) is not None:
|
||||
absolute_pos_embed = state_dict[append_prefix('absolute_pos_embed')]
|
||||
N1, L, C1 = absolute_pos_embed.size()
|
||||
N2, C2, H, W = model.backbone.bottom_up.backbone.absolute_pos_embed.size()
|
||||
if N1 != N2 or C1 != C2 or L != H * W:
|
||||
logger.warning("Error in loading absolute_pos_embed, pass")
|
||||
else:
|
||||
state_dict[append_prefix('absolute_pos_embed')] = absolute_pos_embed.view(N2, H, W, C2).permute(0, 3, 1, 2)
|
||||
|
||||
def get_dist_info():
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
rank, _ = get_dist_info()
|
||||
all_keys = list(state_dict.keys())
|
||||
for key in all_keys:
|
||||
if "relative_position_index" in key:
|
||||
state_dict.pop(key)
|
||||
|
||||
if "relative_position_bias_table" in key:
|
||||
rel_pos_bias = state_dict[key]
|
||||
src_num_pos, num_attn_heads = rel_pos_bias.size()
|
||||
if key not in model.state_dict():
|
||||
continue
|
||||
dst_num_pos, _ = model.state_dict()[key].size()
|
||||
dst_patch_shape = model.backbone.bottom_up.backbone.patch_embed.patch_shape
|
||||
if dst_patch_shape[0] != dst_patch_shape[1]:
|
||||
raise NotImplementedError()
|
||||
num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
|
||||
src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
|
||||
dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
|
||||
if src_size != dst_size:
|
||||
if rank == 0:
|
||||
print("Position interpolate for %s from %dx%d to %dx%d" % (
|
||||
key, src_size, src_size, dst_size, dst_size))
|
||||
extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
|
||||
rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
|
||||
|
||||
def geometric_progression(a, r, n):
|
||||
return a * (1.0 - r ** n) / (1.0 - r)
|
||||
|
||||
left, right = 1.01, 1.5
|
||||
while right - left > 1e-6:
|
||||
q = (left + right) / 2.0
|
||||
gp = geometric_progression(1, q, src_size // 2)
|
||||
if gp > dst_size // 2:
|
||||
right = q
|
||||
else:
|
||||
left = q
|
||||
|
||||
# if q > 1.13492:
|
||||
# q = 1.13492
|
||||
|
||||
dis = []
|
||||
cur = 1
|
||||
for i in range(src_size // 2):
|
||||
dis.append(cur)
|
||||
cur += q ** (i + 1)
|
||||
|
||||
r_ids = [-_ for _ in reversed(dis)]
|
||||
|
||||
x = r_ids + [0] + dis
|
||||
y = r_ids + [0] + dis
|
||||
|
||||
t = dst_size // 2.0
|
||||
dx = np.arange(-t, t + 0.1, 1.0)
|
||||
dy = np.arange(-t, t + 0.1, 1.0)
|
||||
if rank == 0:
|
||||
print("x = {}".format(x))
|
||||
print("dx = {}".format(dx))
|
||||
|
||||
all_rel_pos_bias = []
|
||||
|
||||
for i in range(num_attn_heads):
|
||||
z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
|
||||
f = interpolate.interp2d(x, y, z, kind='cubic')
|
||||
all_rel_pos_bias.append(
|
||||
torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
|
||||
|
||||
rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
|
||||
new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
|
||||
state_dict[key] = new_rel_pos_bias
|
||||
|
||||
if append_prefix('pos_embed') in state_dict:
|
||||
pos_embed_checkpoint = state_dict[append_prefix('pos_embed')]
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.backbone.bottom_up.backbone.patch_embed.num_patches
|
||||
num_extra_tokens = model.backbone.bottom_up.backbone.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
# new_size = int(num_patches ** 0.5)
|
||||
new_size_w = model.backbone.bottom_up.backbone.patch_embed.num_patches_w
|
||||
new_size_h = model.backbone.bottom_up.backbone.patch_embed.num_patches_h
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size_h or orig_size != new_size_w:
|
||||
if rank == 0:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size_w, new_size_h))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size_w, new_size_h), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
state_dict[append_prefix('pos_embed')] = new_pos_embed
|
||||
|
||||
# interpolate position bias table if needed
|
||||
relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
|
||||
for table_key in relative_position_bias_table_keys:
|
||||
table_pretrained = state_dict[table_key]
|
||||
if table_key not in model.state_dict():
|
||||
continue
|
||||
table_current = model.state_dict()[table_key]
|
||||
L1, nH1 = table_pretrained.size()
|
||||
L2, nH2 = table_current.size()
|
||||
if nH1 != nH2:
|
||||
logger.warning(f"Error in loading {table_key}, pass")
|
||||
else:
|
||||
if L1 != L2:
|
||||
S1 = int(L1 ** 0.5)
|
||||
S2 = int(L2 ** 0.5)
|
||||
table_pretrained_resized = F.interpolate(
|
||||
table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
|
||||
size=(S2, S2), mode='bicubic')
|
||||
state_dict[table_key] = table_pretrained_resized.view(nH2, L2).permute(1, 0)
|
||||
|
||||
if append_prefix('rel_pos_bias.relative_position_bias_table') in state_dict and \
|
||||
model.backbone.bottom_up.backbone.use_rel_pos_bias and \
|
||||
not model.backbone.bottom_up.backbone.use_shared_rel_pos_bias and \
|
||||
append_prefix('blocks.0.attn.relative_position_bias_table') not in state_dict:
|
||||
logger.info("[BEIT] Expand the shared relative position embedding to each transformer block. ")
|
||||
num_layers = model.backbone.bottom_up.backbone.get_num_layers()
|
||||
rel_pos_bias = state_dict[append_prefix("rel_pos_bias.relative_position_bias_table")]
|
||||
for i in range(num_layers):
|
||||
state_dict["blocks.%d.attn.relative_position_bias_table" % i] = rel_pos_bias.clone()
|
||||
state_dict.pop(append_prefix("rel_pos_bias.relative_position_bias_table"))
|
||||
|
||||
return state_dict
|
||||
|
||||
|
||||
class MyDetectionCheckpointer(DetectionCheckpointer):
|
||||
def _load_model(self, checkpoint: Any) -> _IncompatibleKeys:
|
||||
"""
|
||||
Load weights from a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint (Any): checkpoint contains the weights.
|
||||
|
||||
Returns:
|
||||
``NamedTuple`` with ``missing_keys``, ``unexpected_keys``,
|
||||
and ``incorrect_shapes`` fields:
|
||||
* **missing_keys** is a list of str containing the missing keys
|
||||
* **unexpected_keys** is a list of str containing the unexpected keys
|
||||
* **incorrect_shapes** is a list of (key, shape in checkpoint, shape in model)
|
||||
|
||||
This is just like the return value of
|
||||
:func:`torch.nn.Module.load_state_dict`, but with extra support
|
||||
for ``incorrect_shapes``.
|
||||
"""
|
||||
checkpoint_state_dict = checkpoint.pop("model")
|
||||
self._convert_ndarray_to_tensor(checkpoint_state_dict)
|
||||
|
||||
# if the state_dict comes from a model that was wrapped in a
|
||||
# DataParallel or DistributedDataParallel during serialization,
|
||||
# remove the "module" prefix before performing the matching.
|
||||
_strip_prefix_if_present(checkpoint_state_dict, "module.")
|
||||
|
||||
# workaround https://github.com/pytorch/pytorch/issues/24139
|
||||
model_state_dict = self.model.state_dict()
|
||||
incorrect_shapes = []
|
||||
|
||||
# rename the para in checkpoint_state_dict
|
||||
# some bug here, do not support re load
|
||||
|
||||
checkpoint_state_dict = {
|
||||
append_prefix(k): checkpoint_state_dict[k]
|
||||
for k in checkpoint_state_dict.keys()
|
||||
}
|
||||
|
||||
checkpoint_state_dict = modify_ckpt_state(self.model, checkpoint_state_dict, logger=self.logger)
|
||||
|
||||
for k in list(checkpoint_state_dict.keys()):
|
||||
if k in model_state_dict:
|
||||
model_param = model_state_dict[k]
|
||||
# Allow mismatch for uninitialized parameters
|
||||
if TORCH_VERSION >= (1, 8) and isinstance(
|
||||
model_param, nn.parameter.UninitializedParameter
|
||||
):
|
||||
continue
|
||||
shape_model = tuple(model_param.shape)
|
||||
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
|
||||
if shape_model != shape_checkpoint:
|
||||
|
||||
has_observer_base_classes = (
|
||||
TORCH_VERSION >= (1, 8)
|
||||
and hasattr(quantization, "ObserverBase")
|
||||
and hasattr(quantization, "FakeQuantizeBase")
|
||||
)
|
||||
if has_observer_base_classes:
|
||||
# Handle the special case of quantization per channel observers,
|
||||
# where buffer shape mismatches are expected.
|
||||
def _get_module_for_key(
|
||||
model: torch.nn.Module, key: str
|
||||
) -> torch.nn.Module:
|
||||
# foo.bar.param_or_buffer_name -> [foo, bar]
|
||||
key_parts = key.split(".")[:-1]
|
||||
cur_module = model
|
||||
for key_part in key_parts:
|
||||
cur_module = getattr(cur_module, key_part)
|
||||
return cur_module
|
||||
|
||||
cls_to_skip = (
|
||||
ObserverBase,
|
||||
FakeQuantizeBase,
|
||||
)
|
||||
target_module = _get_module_for_key(self.model, k)
|
||||
if isinstance(target_module, cls_to_skip):
|
||||
# Do not remove modules with expected shape mismatches
|
||||
# them from the state_dict loading. They have special logic
|
||||
# in _load_from_state_dict to handle the mismatches.
|
||||
continue
|
||||
|
||||
incorrect_shapes.append((k, shape_checkpoint, shape_model))
|
||||
checkpoint_state_dict.pop(k)
|
||||
incompatible = self.model.load_state_dict(checkpoint_state_dict, strict=False)
|
||||
return _IncompatibleKeys(
|
||||
missing_keys=incompatible.missing_keys,
|
||||
unexpected_keys=incompatible.unexpected_keys,
|
||||
incorrect_shapes=incorrect_shapes,
|
||||
)
|
||||
@@ -0,0 +1,744 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
|
||||
"""
|
||||
This file contains components with some default boilerplate logic user may need
|
||||
in training / testing. They will not work for everyone, but many users may find them useful.
|
||||
|
||||
The behavior of functions/classes in this file is subject to change,
|
||||
since they are meant to represent the "common default behavior" people need in their projects.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from typing import Optional
|
||||
import torch
|
||||
from fvcore.nn.precise_bn import get_bn_modules
|
||||
from omegaconf import OmegaConf
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
import detectron2.data.transforms as T
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
from detectron2.config import CfgNode, LazyConfig
|
||||
from detectron2.data import (
|
||||
MetadataCatalog,
|
||||
build_detection_test_loader,
|
||||
build_detection_train_loader,
|
||||
)
|
||||
from detectron2.evaluation import (
|
||||
DatasetEvaluator,
|
||||
inference_on_dataset,
|
||||
print_csv_format,
|
||||
verify_results,
|
||||
)
|
||||
from detectron2.modeling import build_model
|
||||
from detectron2.solver import build_lr_scheduler, build_optimizer
|
||||
from detectron2.utils import comm
|
||||
from detectron2.utils.collect_env import collect_env_info
|
||||
from detectron2.utils.env import seed_all_rng
|
||||
from detectron2.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
|
||||
from detectron2.utils.file_io import PathManager
|
||||
from detectron2.utils.logger import setup_logger
|
||||
|
||||
from detectron2.engine import hooks
|
||||
from detectron2.engine.train_loop import AMPTrainer, SimpleTrainer, TrainerBase
|
||||
|
||||
from .mycheckpointer import MyDetectionCheckpointer
|
||||
from typing import Any, Dict, List, Set
|
||||
import itertools
|
||||
from detectron2.solver.build import maybe_add_gradient_clipping
|
||||
from .dataset_mapper import DetrDatasetMapper
|
||||
from .icdar_evaluation import ICDAREvaluator
|
||||
from detectron2.evaluation import COCOEvaluator
|
||||
|
||||
__all__ = [
|
||||
"create_ddp_model",
|
||||
"default_argument_parser",
|
||||
"default_setup",
|
||||
"default_writers",
|
||||
"DefaultPredictor",
|
||||
"MyTrainer",
|
||||
]
|
||||
|
||||
|
||||
def create_ddp_model(model, *, fp16_compression=False, **kwargs):
|
||||
"""
|
||||
Create a DistributedDataParallel model if there are >1 processes.
|
||||
|
||||
Args:
|
||||
model: a torch.nn.Module
|
||||
fp16_compression: add fp16 compression hooks to the ddp object.
|
||||
See more at https://pytorch.org/docs/stable/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook
|
||||
kwargs: other arguments of :module:`torch.nn.parallel.DistributedDataParallel`.
|
||||
""" # noqa
|
||||
if comm.get_world_size() == 1:
|
||||
return model
|
||||
if "device_ids" not in kwargs:
|
||||
kwargs["device_ids"] = [comm.get_local_rank()]
|
||||
ddp = DistributedDataParallel(model, **kwargs)
|
||||
if fp16_compression:
|
||||
from torch.distributed.algorithms.ddp_comm_hooks import default as comm_hooks
|
||||
|
||||
ddp.register_comm_hook(state=None, hook=comm_hooks.fp16_compress_hook)
|
||||
return ddp
|
||||
|
||||
|
||||
def default_argument_parser(epilog=None):
|
||||
"""
|
||||
Create a parser with some common arguments used by detectron2 users.
|
||||
|
||||
Args:
|
||||
epilog (str): epilog passed to ArgumentParser describing the usage.
|
||||
|
||||
Returns:
|
||||
argparse.ArgumentParser:
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
epilog=epilog
|
||||
or f"""
|
||||
Examples:
|
||||
|
||||
Run on single machine:
|
||||
$ {sys.argv[0]} --num-gpus 8 --config-file cfg.yaml
|
||||
|
||||
Change some config options:
|
||||
$ {sys.argv[0]} --config-file cfg.yaml MODEL.WEIGHTS /path/to/weight.pth SOLVER.BASE_LR 0.001
|
||||
|
||||
Run on multiple machines:
|
||||
(machine0)$ {sys.argv[0]} --machine-rank 0 --num-machines 2 --dist-url <URL> [--other-flags]
|
||||
(machine1)$ {sys.argv[0]} --machine-rank 1 --num-machines 2 --dist-url <URL> [--other-flags]
|
||||
""",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
|
||||
parser.add_argument(
|
||||
"--resume",
|
||||
action="store_true",
|
||||
help="Whether to attempt to resume from the checkpoint directory. "
|
||||
"See documentation of `MyTrainer.resume_or_load()` for what it means.",
|
||||
)
|
||||
parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
|
||||
parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")
|
||||
parser.add_argument("--num-machines", type=int, default=1, help="total number of machines")
|
||||
parser.add_argument(
|
||||
"--machine-rank", type=int, default=0, help="the rank of this machine (unique per machine)"
|
||||
)
|
||||
|
||||
# PyTorch still may leave orphan processes in multi-gpu training.
|
||||
# Therefore we use a deterministic way to obtain port,
|
||||
# so that users are aware of orphan processes by seeing the port occupied.
|
||||
port = 2 ** 15 + 2 ** 14 + hash(os.getuid() if sys.platform != "win32" else 1) % 2 ** 14
|
||||
parser.add_argument(
|
||||
"--dist-url",
|
||||
default="tcp://127.0.0.1:{}".format(port),
|
||||
help="initialization URL for pytorch distributed backend. See "
|
||||
"https://pytorch.org/docs/stable/distributed.html for details.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"opts",
|
||||
help="""
|
||||
Modify config options at the end of the command. For Yacs configs, use
|
||||
space-separated "PATH.KEY VALUE" pairs.
|
||||
For python-based LazyConfig, use "path.key=value".
|
||||
""".strip(),
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def _try_get_key(cfg, *keys, default=None):
|
||||
"""
|
||||
Try select keys from cfg until the first key that exists. Otherwise return default.
|
||||
"""
|
||||
if isinstance(cfg, CfgNode):
|
||||
cfg = OmegaConf.create(cfg.dump())
|
||||
for k in keys:
|
||||
none = object()
|
||||
p = OmegaConf.select(cfg, k, default=none)
|
||||
if p is not none:
|
||||
return p
|
||||
return default
|
||||
|
||||
|
||||
def _highlight(code, filename):
|
||||
try:
|
||||
import pygments
|
||||
except ImportError:
|
||||
return code
|
||||
|
||||
from pygments.lexers import Python3Lexer, YamlLexer
|
||||
from pygments.formatters import Terminal256Formatter
|
||||
|
||||
lexer = Python3Lexer() if filename.endswith(".py") else YamlLexer()
|
||||
code = pygments.highlight(code, lexer, Terminal256Formatter(style="monokai"))
|
||||
return code
|
||||
|
||||
|
||||
def default_setup(cfg, args):
|
||||
"""
|
||||
Perform some basic common setups at the beginning of a job, including:
|
||||
|
||||
1. Set up the detectron2 logger
|
||||
2. Log basic information about environment, cmdline arguments, and config
|
||||
3. Backup the config to the output directory
|
||||
|
||||
Args:
|
||||
cfg (CfgNode or omegaconf.DictConfig): the full config to be used
|
||||
args (argparse.NameSpace): the command line arguments to be logged
|
||||
"""
|
||||
output_dir = _try_get_key(cfg, "OUTPUT_DIR", "output_dir", "train.output_dir")
|
||||
if comm.is_main_process() and output_dir:
|
||||
PathManager.mkdirs(output_dir)
|
||||
|
||||
rank = comm.get_rank()
|
||||
setup_logger(output_dir, distributed_rank=rank, name="fvcore")
|
||||
logger = setup_logger(output_dir, distributed_rank=rank)
|
||||
|
||||
logger.info("Rank of current process: {}. World size: {}".format(rank, comm.get_world_size()))
|
||||
logger.info("Environment info:\n" + collect_env_info())
|
||||
|
||||
logger.info("Command line arguments: " + str(args))
|
||||
if hasattr(args, "config_file") and args.config_file != "":
|
||||
logger.info(
|
||||
"Contents of args.config_file={}:\n{}".format(
|
||||
args.config_file,
|
||||
_highlight(PathManager.open(args.config_file, "r").read(), args.config_file),
|
||||
)
|
||||
)
|
||||
|
||||
if comm.is_main_process() and output_dir:
|
||||
# Note: some of our scripts may expect the existence of
|
||||
# config.yaml in output directory
|
||||
path = os.path.join(output_dir, "config.yaml")
|
||||
if isinstance(cfg, CfgNode):
|
||||
logger.info("Running with full config:\n{}".format(_highlight(cfg.dump(), ".yaml")))
|
||||
with PathManager.open(path, "w") as f:
|
||||
f.write(cfg.dump())
|
||||
else:
|
||||
LazyConfig.save(cfg, path)
|
||||
logger.info("Full config saved to {}".format(path))
|
||||
|
||||
# make sure each worker has a different, yet deterministic seed if specified
|
||||
seed = _try_get_key(cfg, "SEED", "train.seed", default=-1)
|
||||
seed_all_rng(None if seed < 0 else seed + rank)
|
||||
|
||||
# cudnn benchmark has large overhead. It shouldn't be used considering the small size of
|
||||
# typical validation set.
|
||||
if not (hasattr(args, "eval_only") and args.eval_only):
|
||||
torch.backends.cudnn.benchmark = _try_get_key(
|
||||
cfg, "CUDNN_BENCHMARK", "train.cudnn_benchmark", default=False
|
||||
)
|
||||
|
||||
|
||||
def default_writers(output_dir: str, max_iter: Optional[int] = None):
|
||||
"""
|
||||
Build a list of :class:`EventWriter` to be used.
|
||||
It now consists of a :class:`CommonMetricPrinter`,
|
||||
:class:`TensorboardXWriter` and :class:`JSONWriter`.
|
||||
|
||||
Args:
|
||||
output_dir: directory to store JSON metrics and tensorboard events
|
||||
max_iter: the total number of iterations
|
||||
|
||||
Returns:
|
||||
list[EventWriter]: a list of :class:`EventWriter` objects.
|
||||
"""
|
||||
PathManager.mkdirs(output_dir)
|
||||
return [
|
||||
# It may not always print what you want to see, since it prints "common" metrics only.
|
||||
CommonMetricPrinter(max_iter),
|
||||
JSONWriter(os.path.join(output_dir, "metrics.json")),
|
||||
TensorboardXWriter(output_dir),
|
||||
]
|
||||
|
||||
|
||||
class DefaultPredictor:
|
||||
"""
|
||||
Create a simple end-to-end predictor with the given config that runs on
|
||||
single device for a single input image.
|
||||
|
||||
Compared to using the model directly, this class does the following additions:
|
||||
|
||||
1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
|
||||
2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
|
||||
3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
|
||||
4. Take one input image and produce a single output, instead of a batch.
|
||||
|
||||
This is meant for simple demo purposes, so it does the above steps automatically.
|
||||
This is not meant for benchmarks or running complicated inference logic.
|
||||
If you'd like to do anything more complicated, please refer to its source code as
|
||||
examples to build and use the model manually.
|
||||
|
||||
Attributes:
|
||||
metadata (Metadata): the metadata of the underlying dataset, obtained from
|
||||
cfg.DATASETS.TEST.
|
||||
|
||||
Examples:
|
||||
::
|
||||
pred = DefaultPredictor(cfg)
|
||||
inputs = cv2.imread("input.jpg")
|
||||
outputs = pred(inputs)
|
||||
"""
|
||||
|
||||
def __init__(self, cfg):
|
||||
self.cfg = cfg.clone() # cfg can be modified by model
|
||||
self.model = build_model(self.cfg)
|
||||
self.model.eval()
|
||||
if len(cfg.DATASETS.TEST):
|
||||
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
|
||||
|
||||
checkpointer = DetectionCheckpointer(self.model)
|
||||
checkpointer.load(cfg.MODEL.WEIGHTS)
|
||||
|
||||
self.aug = T.ResizeShortestEdge(
|
||||
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
|
||||
)
|
||||
|
||||
self.input_format = cfg.INPUT.FORMAT
|
||||
assert self.input_format in ["RGB", "BGR"], self.input_format
|
||||
|
||||
def __call__(self, original_image):
|
||||
"""
|
||||
Args:
|
||||
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
|
||||
|
||||
Returns:
|
||||
predictions (dict):
|
||||
the output of the model for one image only.
|
||||
See :doc:`/tutorials/models` for details about the format.
|
||||
"""
|
||||
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
|
||||
# Apply pre-processing to image.
|
||||
if self.input_format == "RGB":
|
||||
# whether the model expects BGR inputs or RGB
|
||||
original_image = original_image[:, :, ::-1]
|
||||
height, width = original_image.shape[:2]
|
||||
image = self.aug.get_transform(original_image).apply_image(original_image)
|
||||
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
|
||||
|
||||
inputs = {"image": image, "height": height, "width": width}
|
||||
predictions = self.model([inputs])[0]
|
||||
return predictions
|
||||
|
||||
|
||||
class MyTrainer(TrainerBase):
|
||||
"""
|
||||
A trainer with default training logic. It does the following:
|
||||
|
||||
1. Create a :class:`SimpleTrainer` using model, optimizer, dataloader
|
||||
defined by the given config. Create a LR scheduler defined by the config.
|
||||
2. Load the last checkpoint or `cfg.MODEL.WEIGHTS`, if exists, when
|
||||
`resume_or_load` is called.
|
||||
3. Register a few common hooks defined by the config.
|
||||
|
||||
It is created to simplify the **standard model training workflow** and reduce code boilerplate
|
||||
for users who only need the standard training workflow, with standard features.
|
||||
It means this class makes *many assumptions* about your training logic that
|
||||
may easily become invalid in a new research. In fact, any assumptions beyond those made in the
|
||||
:class:`SimpleTrainer` are too much for research.
|
||||
|
||||
The code of this class has been annotated about restrictive assumptions it makes.
|
||||
When they do not work for you, you're encouraged to:
|
||||
|
||||
1. Overwrite methods of this class, OR:
|
||||
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
|
||||
nothing else. You can then add your own hooks if needed. OR:
|
||||
3. Write your own training loop similar to `tools/plain_train_net.py`.
|
||||
|
||||
See the :doc:`/tutorials/training` tutorials for more details.
|
||||
|
||||
Note that the behavior of this class, like other functions/classes in
|
||||
this file, is not stable, since it is meant to represent the "common default behavior".
|
||||
It is only guaranteed to work well with the standard models and training workflow in detectron2.
|
||||
To obtain more stable behavior, write your own training logic with other public APIs.
|
||||
|
||||
Examples:
|
||||
::
|
||||
trainer = MyTrainer(cfg)
|
||||
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
|
||||
trainer.train()
|
||||
|
||||
Attributes:
|
||||
scheduler:
|
||||
checkpointer (DetectionCheckpointer):
|
||||
cfg (CfgNode):
|
||||
"""
|
||||
|
||||
def __init__(self, cfg):
|
||||
"""
|
||||
Args:
|
||||
cfg (CfgNode):
|
||||
"""
|
||||
super().__init__()
|
||||
logger = logging.getLogger("detectron2")
|
||||
if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2
|
||||
setup_logger()
|
||||
cfg = MyTrainer.auto_scale_workers(cfg, comm.get_world_size())
|
||||
|
||||
self.cfg = cfg
|
||||
|
||||
# Assume these objects must be constructed in this order.
|
||||
model = self.build_model(cfg)
|
||||
optimizer = self.build_optimizer(cfg, model)
|
||||
data_loader = self.build_train_loader(cfg)
|
||||
|
||||
model = create_ddp_model(model, broadcast_buffers=False)
|
||||
self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
|
||||
model, data_loader, optimizer
|
||||
)
|
||||
|
||||
self.scheduler = self.build_lr_scheduler(cfg, optimizer)
|
||||
self.checkpointer = MyDetectionCheckpointer(
|
||||
# Assume you want to save checkpoints together with logs/statistics
|
||||
model,
|
||||
cfg.OUTPUT_DIR,
|
||||
trainer=weakref.proxy(self),
|
||||
)
|
||||
self.start_iter = 0
|
||||
self.max_iter = cfg.SOLVER.MAX_ITER
|
||||
self.cfg = cfg
|
||||
|
||||
self.register_hooks(self.build_hooks())
|
||||
|
||||
def resume_or_load(self, resume=True):
|
||||
"""
|
||||
If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
|
||||
a `last_checkpoint` file), resume from the file. Resuming means loading all
|
||||
available states (eg. optimizer and scheduler) and update iteration counter
|
||||
from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
|
||||
|
||||
Otherwise, this is considered as an independent training. The method will load model
|
||||
weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
|
||||
from iteration 0.
|
||||
|
||||
Args:
|
||||
resume (bool): whether to do resume or not
|
||||
"""
|
||||
self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
|
||||
if resume and self.checkpointer.has_checkpoint():
|
||||
# The checkpoint stores the training iteration that just finished, thus we start
|
||||
# at the next iteration
|
||||
self.start_iter = self.iter + 1
|
||||
|
||||
def build_hooks(self):
|
||||
"""
|
||||
Build a list of default hooks, including timing, evaluation,
|
||||
checkpointing, lr scheduling, precise BN, writing events.
|
||||
|
||||
Returns:
|
||||
list[HookBase]:
|
||||
"""
|
||||
cfg = self.cfg.clone()
|
||||
cfg.defrost()
|
||||
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
|
||||
|
||||
ret = [
|
||||
hooks.IterationTimer(),
|
||||
hooks.LRScheduler(),
|
||||
hooks.PreciseBN(
|
||||
# Run at the same freq as (but before) evaluation.
|
||||
cfg.TEST.EVAL_PERIOD,
|
||||
self.model,
|
||||
# Build a new data loader to not affect training
|
||||
self.build_train_loader(cfg),
|
||||
cfg.TEST.PRECISE_BN.NUM_ITER,
|
||||
)
|
||||
if cfg.TEST.PRECISE_BN.ENABLED and get_bn_modules(self.model)
|
||||
else None,
|
||||
]
|
||||
|
||||
# Do PreciseBN before checkpointer, because it updates the model and need to
|
||||
# be saved by checkpointer.
|
||||
# This is not always the best: if checkpointing has a different frequency,
|
||||
# some checkpoints may have more precise statistics than others.
|
||||
if comm.is_main_process():
|
||||
ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
|
||||
|
||||
def test_and_save_results():
|
||||
self._last_eval_results = self.test(self.cfg, self.model)
|
||||
return self._last_eval_results
|
||||
|
||||
# Do evaluation after checkpointer, because then if it fails,
|
||||
# we can use the saved checkpoint to debug.
|
||||
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
|
||||
|
||||
if comm.is_main_process():
|
||||
# Here the default print/log frequency of each writer is used.
|
||||
# run writers in the end, so that evaluation metrics are written
|
||||
ret.append(hooks.PeriodicWriter(self.build_writers(), period=20))
|
||||
return ret
|
||||
|
||||
def build_writers(self):
|
||||
"""
|
||||
Build a list of writers to be used using :func:`default_writers()`.
|
||||
If you'd like a different list of writers, you can overwrite it in
|
||||
your trainer.
|
||||
|
||||
Returns:
|
||||
list[EventWriter]: a list of :class:`EventWriter` objects.
|
||||
"""
|
||||
return default_writers(self.cfg.OUTPUT_DIR, self.max_iter)
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
Run training.
|
||||
|
||||
Returns:
|
||||
OrderedDict of results, if evaluation is enabled. Otherwise None.
|
||||
"""
|
||||
super().train(self.start_iter, self.max_iter)
|
||||
if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
|
||||
assert hasattr(
|
||||
self, "_last_eval_results"
|
||||
), "No evaluation results obtained during training!"
|
||||
verify_results(self.cfg, self._last_eval_results)
|
||||
return self._last_eval_results
|
||||
|
||||
def run_step(self):
|
||||
self._trainer.iter = self.iter
|
||||
self._trainer.run_step()
|
||||
|
||||
@classmethod
|
||||
def build_model(cls, cfg):
|
||||
"""
|
||||
Returns:
|
||||
torch.nn.Module:
|
||||
|
||||
It now calls :func:`detectron2.modeling.build_model`.
|
||||
Overwrite it if you'd like a different model.
|
||||
"""
|
||||
model = build_model(cfg)
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info("Model:\n{}".format(model))
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def build_optimizer(cls, cfg, model):
|
||||
params: List[Dict[str, Any]] = []
|
||||
memo: Set[torch.nn.parameter.Parameter] = set()
|
||||
for key, value in model.named_parameters(recurse=True):
|
||||
if not value.requires_grad:
|
||||
continue
|
||||
# Avoid duplicating parameters
|
||||
if value in memo:
|
||||
continue
|
||||
memo.add(value)
|
||||
lr = cfg.SOLVER.BASE_LR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY
|
||||
if "backbone" in key:
|
||||
lr = lr * cfg.SOLVER.BACKBONE_MULTIPLIER
|
||||
params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
|
||||
|
||||
def maybe_add_full_model_gradient_clipping(optim): # optim: the optimizer class
|
||||
# detectron2 doesn't have full model gradient clipping now
|
||||
clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
|
||||
enable = (
|
||||
cfg.SOLVER.CLIP_GRADIENTS.ENABLED
|
||||
and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
|
||||
and clip_norm_val > 0.0
|
||||
)
|
||||
|
||||
class FullModelGradientClippingOptimizer(optim):
|
||||
def step(self, closure=None):
|
||||
all_params = itertools.chain(*[x["params"] for x in self.param_groups])
|
||||
torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
|
||||
super().step(closure=closure)
|
||||
|
||||
return FullModelGradientClippingOptimizer if enable else optim
|
||||
|
||||
optimizer_type = cfg.SOLVER.OPTIMIZER
|
||||
if optimizer_type == "SGD":
|
||||
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
|
||||
params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
|
||||
)
|
||||
elif optimizer_type == "ADAMW":
|
||||
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
|
||||
params, cfg.SOLVER.BASE_LR
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"no optimizer type {optimizer_type}")
|
||||
if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
|
||||
optimizer = maybe_add_gradient_clipping(cfg, optimizer)
|
||||
return optimizer
|
||||
|
||||
@classmethod
|
||||
def build_lr_scheduler(cls, cfg, optimizer):
|
||||
"""
|
||||
It now calls :func:`detectron2.solver.build_lr_scheduler`.
|
||||
Overwrite it if you'd like a different scheduler.
|
||||
"""
|
||||
return build_lr_scheduler(cfg, optimizer)
|
||||
|
||||
@classmethod
|
||||
def build_train_loader(cls, cfg):
|
||||
if cfg.AUG.DETR:
|
||||
mapper = DetrDatasetMapper(cfg, is_train=True)
|
||||
else:
|
||||
mapper = None
|
||||
return build_detection_train_loader(cfg, mapper=mapper)
|
||||
|
||||
@classmethod
|
||||
def build_test_loader(cls, cfg, dataset_name):
|
||||
"""
|
||||
Returns:
|
||||
iterable
|
||||
|
||||
It now calls :func:`detectron2.data.build_detection_test_loader`.
|
||||
Overwrite it if you'd like a different data loader.
|
||||
"""
|
||||
return build_detection_test_loader(cfg, dataset_name)
|
||||
|
||||
@classmethod
|
||||
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
|
||||
if output_folder is None:
|
||||
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
|
||||
if 'icdar' not in dataset_name:
|
||||
return COCOEvaluator(dataset_name, output_dir=output_folder)
|
||||
else:
|
||||
return ICDAREvaluator(dataset_name, output_dir=output_folder)
|
||||
|
||||
@classmethod
|
||||
def test(cls, cfg, model, evaluators=None):
|
||||
"""
|
||||
Evaluate the given model. The given model is expected to already contain
|
||||
weights to evaluate.
|
||||
|
||||
Args:
|
||||
cfg (CfgNode):
|
||||
model (nn.Module):
|
||||
evaluators (list[DatasetEvaluator] or None): if None, will call
|
||||
:meth:`build_evaluator`. Otherwise, must have the same length as
|
||||
``cfg.DATASETS.TEST``.
|
||||
|
||||
Returns:
|
||||
dict: a dict of result metrics
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
if isinstance(evaluators, DatasetEvaluator):
|
||||
evaluators = [evaluators]
|
||||
if evaluators is not None:
|
||||
assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
|
||||
len(cfg.DATASETS.TEST), len(evaluators)
|
||||
)
|
||||
|
||||
results = OrderedDict()
|
||||
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
|
||||
data_loader = cls.build_test_loader(cfg, dataset_name)
|
||||
# When evaluators are passed in as arguments,
|
||||
# implicitly assume that evaluators can be created before data_loader.
|
||||
if evaluators is not None:
|
||||
evaluator = evaluators[idx]
|
||||
else:
|
||||
try:
|
||||
evaluator = cls.build_evaluator(cfg, dataset_name)
|
||||
except NotImplementedError:
|
||||
logger.warn(
|
||||
"No evaluator found. Use `MyTrainer.test(evaluators=)`, "
|
||||
"or implement its `build_evaluator` method."
|
||||
)
|
||||
results[dataset_name] = {}
|
||||
continue
|
||||
results_i = inference_on_dataset(model, data_loader, evaluator)
|
||||
results[dataset_name] = results_i
|
||||
if comm.is_main_process():
|
||||
assert isinstance(
|
||||
results_i, dict
|
||||
), "Evaluator must return a dict on the main process. Got {} instead.".format(
|
||||
results_i
|
||||
)
|
||||
logger.info("Evaluation results for {} in csv format:".format(dataset_name))
|
||||
print_csv_format(results_i)
|
||||
|
||||
if len(results) == 1:
|
||||
results = list(results.values())[0]
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def auto_scale_workers(cfg, num_workers: int):
|
||||
"""
|
||||
When the config is defined for certain number of workers (according to
|
||||
``cfg.SOLVER.REFERENCE_WORLD_SIZE``) that's different from the number of
|
||||
workers currently in use, returns a new cfg where the total batch size
|
||||
is scaled so that the per-GPU batch size stays the same as the
|
||||
original ``IMS_PER_BATCH // REFERENCE_WORLD_SIZE``.
|
||||
|
||||
Other config options are also scaled accordingly:
|
||||
* training steps and warmup steps are scaled inverse proportionally.
|
||||
* learning rate are scaled proportionally, following :paper:`ImageNet in 1h`.
|
||||
|
||||
For example, with the original config like the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
IMS_PER_BATCH: 16
|
||||
BASE_LR: 0.1
|
||||
REFERENCE_WORLD_SIZE: 8
|
||||
MAX_ITER: 5000
|
||||
STEPS: (4000,)
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
|
||||
When this config is used on 16 GPUs instead of the reference number 8,
|
||||
calling this method will return a new config with:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
IMS_PER_BATCH: 32
|
||||
BASE_LR: 0.2
|
||||
REFERENCE_WORLD_SIZE: 16
|
||||
MAX_ITER: 2500
|
||||
STEPS: (2000,)
|
||||
CHECKPOINT_PERIOD: 500
|
||||
|
||||
Note that both the original config and this new config can be trained on 16 GPUs.
|
||||
It's up to user whether to enable this feature (by setting ``REFERENCE_WORLD_SIZE``).
|
||||
|
||||
Returns:
|
||||
CfgNode: a new config. Same as original if ``cfg.SOLVER.REFERENCE_WORLD_SIZE==0``.
|
||||
"""
|
||||
old_world_size = cfg.SOLVER.REFERENCE_WORLD_SIZE
|
||||
if old_world_size == 0 or old_world_size == num_workers:
|
||||
return cfg
|
||||
cfg = cfg.clone()
|
||||
frozen = cfg.is_frozen()
|
||||
cfg.defrost()
|
||||
|
||||
assert (
|
||||
cfg.SOLVER.IMS_PER_BATCH % old_world_size == 0
|
||||
), "Invalid REFERENCE_WORLD_SIZE in config!"
|
||||
scale = num_workers / old_world_size
|
||||
bs = cfg.SOLVER.IMS_PER_BATCH = int(round(cfg.SOLVER.IMS_PER_BATCH * scale))
|
||||
lr = cfg.SOLVER.BASE_LR = cfg.SOLVER.BASE_LR * scale
|
||||
max_iter = cfg.SOLVER.MAX_ITER = int(round(cfg.SOLVER.MAX_ITER / scale))
|
||||
warmup_iter = cfg.SOLVER.WARMUP_ITERS = int(round(cfg.SOLVER.WARMUP_ITERS / scale))
|
||||
cfg.SOLVER.STEPS = tuple(int(round(s / scale)) for s in cfg.SOLVER.STEPS)
|
||||
cfg.TEST.EVAL_PERIOD = int(round(cfg.TEST.EVAL_PERIOD / scale))
|
||||
cfg.SOLVER.CHECKPOINT_PERIOD = int(round(cfg.SOLVER.CHECKPOINT_PERIOD / scale))
|
||||
cfg.SOLVER.REFERENCE_WORLD_SIZE = num_workers # maintain invariant
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(
|
||||
f"Auto-scaling the config to batch_size={bs}, learning_rate={lr}, "
|
||||
f"max_iter={max_iter}, warmup={warmup_iter}."
|
||||
)
|
||||
|
||||
if frozen:
|
||||
cfg.freeze()
|
||||
return cfg
|
||||
|
||||
|
||||
# Access basic attributes from the underlying trainer
|
||||
for _attr in ["model", "data_loader", "optimizer"]:
|
||||
setattr(
|
||||
MyTrainer,
|
||||
_attr,
|
||||
property(
|
||||
# getter
|
||||
lambda self, x=_attr: getattr(self._trainer, x),
|
||||
# setter
|
||||
lambda self, value, x=_attr: setattr(self._trainer, x, value),
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
from .evaluate import calc_table_score
|
||||
@@ -0,0 +1,469 @@
|
||||
"""
|
||||
Data structures used by the evaluation process.
|
||||
Yu Fang - March 2019
|
||||
"""
|
||||
|
||||
from collections import Iterable
|
||||
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
|
||||
# helper functions
|
||||
def flatten(lis):
|
||||
for item in lis:
|
||||
if isinstance(item, Iterable) and not isinstance(item, str):
|
||||
for x in flatten(item):
|
||||
yield x
|
||||
else:
|
||||
yield item
|
||||
|
||||
# derived from https://blog.csdn.net/u012433049/article/details/82909484
|
||||
def compute_poly_iou(list1, list2):
|
||||
a1 = np.array(list1, dtype=int).reshape(-1, 2)
|
||||
poly1 = Polygon(a1)
|
||||
poly1_clean = poly1.buffer(0)
|
||||
|
||||
a2 = np.array(list2, dtype=int).reshape(-1, 2)
|
||||
poly2 = Polygon(a2)
|
||||
poly2_clean = poly2.buffer(0)
|
||||
|
||||
try:
|
||||
# iou = poly1.intersection(poly2).area / poly1.union(poly2).area
|
||||
iou = poly1_clean.intersection(poly2_clean).area / poly1_clean.union(poly2_clean).area
|
||||
except ZeroDivisionError:
|
||||
iou = 0
|
||||
return iou
|
||||
|
||||
|
||||
class Cell(object):
|
||||
# @:param start_row : start row index of the Cell
|
||||
# @:param start_col : start column index of the Cell
|
||||
# @:param end-row : end row index of the Cell
|
||||
# @:param end-col : end column index of the Cell
|
||||
# @:param cell_box: bounding-box of the Cell (coordinates are saved as a string)
|
||||
# @:param content_box: bounding-box of the text content within Cell (unused variable)
|
||||
# @:param cell_id: unique id of the Cell
|
||||
|
||||
def __init__(self, table_id, start_row, start_col, cell_box, end_row, end_col, content_box=""):
|
||||
self._start_row = int(start_row)
|
||||
self._start_col = int(start_col)
|
||||
self._cell_box = cell_box
|
||||
self._content_box = content_box
|
||||
self._table_id = table_id # the table_id this cell belongs to
|
||||
# self._cell_name = cell_id # specify the cell using passed-in cell_id
|
||||
self._cell_id = id(self)
|
||||
# self._region = region
|
||||
|
||||
# check for end-row and end-col special case
|
||||
if end_row == -1:
|
||||
self._end_row = self.start_row
|
||||
else:
|
||||
self._end_row = int(end_row)
|
||||
if end_col == -1:
|
||||
self._end_col = self._start_col
|
||||
else:
|
||||
self._end_col = int(end_col)
|
||||
|
||||
@property
|
||||
def start_row(self):
|
||||
return self._start_row
|
||||
|
||||
@property
|
||||
def start_col(self):
|
||||
return self._start_col
|
||||
|
||||
@property
|
||||
def end_row(self):
|
||||
return self._end_row
|
||||
|
||||
@property
|
||||
def end_col(self):
|
||||
return self._end_col
|
||||
|
||||
@property
|
||||
def cell_box(self):
|
||||
return self._cell_box
|
||||
|
||||
@property
|
||||
def content_box(self):
|
||||
return self._content_box
|
||||
|
||||
@property
|
||||
def cell_id(self):
|
||||
return self._cell_id
|
||||
|
||||
@property
|
||||
def table_id(self):
|
||||
return self._table_id
|
||||
|
||||
def __str__(self):
|
||||
return "CELL row=[%d, %d] col=[%d, %d] (coords=%s)" %(self.start_row, self.end_row
|
||||
, self.start_col, self.end_col
|
||||
, self.cell_box)
|
||||
|
||||
# return the IoU value of two cell blocks
|
||||
def compute_cell_iou(self, another_cell):
|
||||
cell_box_1_temp = []
|
||||
for el in self.cell_box.split():
|
||||
cell_box_1_temp.append((el.split(",")))
|
||||
cell_box_1 = list(flatten(cell_box_1_temp))
|
||||
cell_box_1 = [int(x) for x in cell_box_1]
|
||||
|
||||
cell_box_2_temp = []
|
||||
for el in another_cell.cell_box.split():
|
||||
cell_box_2_temp.append((el.split(",")))
|
||||
cell_box_2 = list(flatten(cell_box_2_temp))
|
||||
cell_box_2 = [int(x) for x in cell_box_2]
|
||||
|
||||
return compute_poly_iou(cell_box_1, cell_box_2)
|
||||
|
||||
# check if the two cell object denotes same cell area in table
|
||||
def check_same(self, another_cell):
|
||||
return self._start_row == another_cell.start_row and self._end_row == another_cell.end_row and \
|
||||
self._start_col == another_cell.start_col and self._end_col == another_cell.end_col
|
||||
|
||||
|
||||
# Note: currently save the relation with two cell object involved,
|
||||
# can be replaced by cell_id in follow-up memory clean up
|
||||
class AdjRelation:
|
||||
|
||||
DIR_HORIZ = 1
|
||||
DIR_VERT = 2
|
||||
|
||||
def __init__(self, fromText, toText, direction):
|
||||
# @param: fromText, toText are Cell objects （may be changed to cell-ID for further development）
|
||||
self._fromText = fromText
|
||||
self._toText = toText
|
||||
self._direction = direction
|
||||
|
||||
@property
|
||||
def fromText(self):
|
||||
return self._fromText
|
||||
|
||||
@property
|
||||
def toText(self):
|
||||
return self._toText
|
||||
|
||||
@property
|
||||
def direction(self):
|
||||
return self._direction
|
||||
|
||||
def __str__(self):
|
||||
if self.direction == self.DIR_VERT:
|
||||
dir = "vertical"
|
||||
else:
|
||||
dir = "horizontal"
|
||||
return 'ADJ_RELATION: ' + str(self._fromText) + ' ' + str(self._toText) + ' ' + dir
|
||||
|
||||
def isEqual(self, otherRelation):
|
||||
return self.fromText.cell_id == otherRelation.fromText.cell_id and \
|
||||
self.toText.cell_id == otherRelation.toText.cell_id and self.direction == otherRelation.direction
|
||||
|
||||
|
||||
class Table:
|
||||
|
||||
def __init__(self, tableNode):
|
||||
self._root = tableNode
|
||||
self._id = id(self)
|
||||
self._table_coords = ""
|
||||
self._maxRow = 0 # PS: indexing from 0
|
||||
self._maxCol = 0
|
||||
self._cells = [] # save a table as list of <Cell>s
|
||||
self.adj_relations = [] # save the adj_relations for the table
|
||||
self.parsed = False
|
||||
self.found = False # check if the find_adj_relations() has been called once
|
||||
|
||||
self.parse_table()
|
||||
|
||||
def __str__(self):
|
||||
return "TABLE object - {} row x {} col".format(self._maxRow+1, self._maxCol+1)
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self._id
|
||||
|
||||
@property
|
||||
def table_coords(self):
|
||||
return self._table_coords
|
||||
|
||||
@property
|
||||
def table_cells(self):
|
||||
return self._cells
|
||||
|
||||
# parse input xml to cell lists
|
||||
def parse_table(self):
|
||||
# get the table bbox
|
||||
self._table_coords = str(self._root.getElementsByTagName("Coords")[0].getAttribute("points"))
|
||||
|
||||
# get info for each cell
|
||||
cells = self._root.getElementsByTagName("cell")
|
||||
max_row = max_col = 0
|
||||
for cell in cells:
|
||||
sr = cell.getAttribute("start-row")
|
||||
sc = cell.getAttribute("start-col")
|
||||
cell_id = cell.getAttribute("id")
|
||||
b_points = str(cell.getElementsByTagName("Coords")[0].getAttribute("points"))
|
||||
# try:
|
||||
# try:
|
||||
# text = cell.getElementsByTagName("content")[0].firstChild.nodeValue
|
||||
# except AttributeError:
|
||||
# text = ""
|
||||
# except IndexError:
|
||||
# text = "initialized cell as no content"
|
||||
er = cell.getAttribute("end-row") if cell.hasAttribute("end-row") else -1
|
||||
ec = cell.getAttribute("end-col") if cell.hasAttribute("end-col") else -1
|
||||
new_cell = Cell(table_id=str(self.id), start_row=sr, start_col=sc, cell_box=b_points,
|
||||
end_row=er, end_col=ec)
|
||||
max_row = max(max_row, int(sr), int(er))
|
||||
max_col = max(max_col, int(sc), int(ec))
|
||||
self._cells.append(new_cell)
|
||||
self._maxCol = max_col
|
||||
self._maxRow = max_row
|
||||
self.parsed = True
|
||||
|
||||
# generate a table-like structure for finding adj_relations
|
||||
def convert_2d(self):
|
||||
table = [[0 for x in range(self._maxCol+1)] for y in range(self._maxRow+1)] # init blank cell with int 0
|
||||
for cell in self._cells:
|
||||
cur_row = cell.start_row
|
||||
while cur_row <= cell.end_row:
|
||||
cur_col = cell.start_col
|
||||
while cur_col <= cell.end_col:
|
||||
temp = table[cur_row][cur_col]
|
||||
if temp == 0:
|
||||
table[cur_row][cur_col] = cell
|
||||
elif type(temp) == list:
|
||||
temp.append(cell)
|
||||
table[cur_row][cur_col] = temp
|
||||
else:
|
||||
table[cur_row][cur_col] = [temp, cell]
|
||||
cur_col += 1
|
||||
cur_row += 1
|
||||
|
||||
return table
|
||||
|
||||
def find_adj_relations(self):
|
||||
if self.found:
|
||||
return self.adj_relations
|
||||
else:
|
||||
# if len(self._cells) == 0:
|
||||
if self.parsed == False:
|
||||
# fix: cases where there's no cell in table?
|
||||
print("table is not parsed for further steps.")
|
||||
self.parse_table()
|
||||
self.find_adj_relations()
|
||||
else:
|
||||
retVal = []
|
||||
tab = self.convert_2d()
|
||||
|
||||
# find horizontal relations
|
||||
for r in range(self._maxRow+1):
|
||||
for c_from in range(self._maxCol):
|
||||
temp_pos = tab[r][c_from]
|
||||
if temp_pos == 0:
|
||||
continue
|
||||
elif type(temp_pos) == list:
|
||||
for cell in temp_pos:
|
||||
c_to = c_from + 1
|
||||
if tab[r][c_to] != 0:
|
||||
# find relation between two adjacent cells
|
||||
if type(tab[r][c_to]) == list:
|
||||
for cell_to in tab[r][c_to]:
|
||||
if cell != cell_to and (not cell.check_same(cell_to)):
|
||||
adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
if cell != tab[r][c_to]:
|
||||
adj_relation = AdjRelation(cell, tab[r][c_to], AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
# find the next non-blank cell, if exists
|
||||
for temp in range(c_from + 1, self._maxCol + 1):
|
||||
if tab[r][temp] != 0:
|
||||
if type(tab[r][temp]) == list:
|
||||
for cell_to in tab[r][temp]:
|
||||
adj_relation = AdjRelation(cell, cell_to,
|
||||
AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
adj_relation = AdjRelation(cell, tab[r][temp],
|
||||
AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
break
|
||||
else:
|
||||
c_to = c_from + 1
|
||||
if tab[r][c_to] != 0:
|
||||
# find relation between two adjacent cells
|
||||
if type(tab[r][c_to]) == list:
|
||||
for cell_to in tab[r][c_to]:
|
||||
if temp_pos != cell_to:
|
||||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
if temp_pos != tab[r][c_to]:
|
||||
adj_relation = AdjRelation(temp_pos, tab[r][c_to], AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
# find the next non-blank cell, if exists
|
||||
for temp in range(c_from + 1, self._maxCol + 1):
|
||||
if tab[r][temp] != 0:
|
||||
if type(tab[r][temp]) == list:
|
||||
for cell_to in tab[r][temp]:
|
||||
adj_relation = AdjRelation(temp_pos, cell_to,
|
||||
AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
adj_relation = AdjRelation(temp_pos, tab[r][temp], AdjRelation.DIR_HORIZ)
|
||||
retVal.append(adj_relation)
|
||||
break
|
||||
|
||||
# find vertical relations
|
||||
for c in range(self._maxCol+1):
|
||||
for r_from in range(self._maxRow):
|
||||
temp_pos = tab[r_from][c]
|
||||
if temp_pos == 0:
|
||||
continue
|
||||
elif type(temp_pos) == list:
|
||||
for cell in temp_pos:
|
||||
r_to = r_from + 1
|
||||
if tab[r_to][c] != 0:
|
||||
# find relation between two adjacent cells
|
||||
if type(tab[r_to][c]) == list:
|
||||
for cell_to in tab[r_to][c]:
|
||||
if cell != cell_to and (not cell.check_same(cell_to)):
|
||||
adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
if cell != tab[r_to][c]:
|
||||
adj_relation = AdjRelation(cell, tab[r_to][c], AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
# find the next non-blank cell, if exists
|
||||
for temp in range(r_from + 1, self._maxRow + 1):
|
||||
if tab[temp][c] != 0:
|
||||
if type(tab[temp][c]) == list:
|
||||
for cell_to in tab[temp][c]:
|
||||
adj_relation = AdjRelation(cell, cell_to,
|
||||
AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
adj_relation = AdjRelation(cell, tab[temp][c],
|
||||
AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
break
|
||||
else:
|
||||
r_to = r_from + 1
|
||||
if tab[r_to][c] != 0:
|
||||
# find relation between two adjacent cells
|
||||
if type(tab[r_to][c]) == list:
|
||||
for cell_to in tab[r_to][c]:
|
||||
if temp_pos != cell_to:
|
||||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
if temp_pos != tab[r_to][c]:
|
||||
adj_relation = AdjRelation(temp_pos, tab[r_to][c], AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
# find the next non-blank cell, if exists
|
||||
for temp in range(r_from + 1, self._maxRow + 1):
|
||||
if tab[temp][c] != 0:
|
||||
if type(tab[temp][c]) == list:
|
||||
for cell_to in tab[temp][c]:
|
||||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
else:
|
||||
adj_relation = AdjRelation(temp_pos, tab[temp][c], AdjRelation.DIR_VERT)
|
||||
retVal.append(adj_relation)
|
||||
break
|
||||
|
||||
# eliminate duplicates
|
||||
repeat = True
|
||||
while repeat:
|
||||
repeat = False
|
||||
duplicates = []
|
||||
|
||||
for ar1 in retVal:
|
||||
for ar2 in retVal:
|
||||
if ar1 != ar2:
|
||||
if ar1.direction == ar2.direction and ar1.fromText == ar2.fromText and\
|
||||
ar1.toText == ar2.toText:
|
||||
duplicates.append(ar2)
|
||||
break
|
||||
else:
|
||||
continue
|
||||
break
|
||||
|
||||
if len(duplicates) > 0:
|
||||
repeat = True
|
||||
retVal.remove(duplicates[0])
|
||||
|
||||
self.found = True
|
||||
self.adj_relations = retVal
|
||||
return self.adj_relations
|
||||
|
||||
# compute the IOU of table, pass-in var is another Table object
|
||||
def compute_table_iou(self, another_table):
|
||||
table_box_1_temp = []
|
||||
for el in self.table_coords.split():
|
||||
table_box_1_temp.append((el.split(",")))
|
||||
table_box_1 = list(flatten(table_box_1_temp))
|
||||
table_box_1 = [int(x) for x in table_box_1]
|
||||
|
||||
table_box_2_temp = []
|
||||
for el in another_table.table_coords.split():
|
||||
table_box_2_temp.append((el.split(",")))
|
||||
table_box_2 = list(flatten(table_box_2_temp))
|
||||
table_box_2 = [int(x) for x in table_box_2]
|
||||
|
||||
return compute_poly_iou(table_box_1, table_box_2)
|
||||
|
||||
# find the cell mapping of tables as dictionary, pass-in var is another table and the desired IOU value
|
||||
def find_cell_mapping(self, target_table, iou_value):
|
||||
mapped_cell = [] # store the matches as tuples - (gt, result) mind the order of table when passing in
|
||||
for cell_1 in self.table_cells:
|
||||
for cell_2 in target_table.table_cells:
|
||||
if cell_1.compute_cell_iou(cell_2) >= iou_value:
|
||||
mapped_cell.append((cell_1, cell_2))
|
||||
break
|
||||
ret = dict(mapped_cell)
|
||||
# print(ret)
|
||||
return ret
|
||||
|
||||
# to print a table cell mapping
|
||||
@classmethod
|
||||
def printCellMapping(cls, dMappedCell):
|
||||
print("-"*25)
|
||||
for cell1, cell2 in dMappedCell.items():
|
||||
print(" ", cell1, " --> ", cell2)
|
||||
|
||||
# to print a table set of adjacency relations
|
||||
@classmethod
|
||||
def printAdjacencyRelationList(cls, lAdjRel, title=""):
|
||||
print("--- %s "%title + "-"*25)
|
||||
for adj in lAdjRel:
|
||||
print(adj)
|
||||
|
||||
|
||||
class ResultStructure:
|
||||
|
||||
def __init__(self, truePos, gtTotal, resTotal):
|
||||
self._truePos = truePos
|
||||
self._gtTotal = gtTotal
|
||||
self._resTotal = resTotal
|
||||
|
||||
@property
|
||||
def truePos(self):
|
||||
return self._truePos
|
||||
|
||||
@property
|
||||
def gtTotal(self):
|
||||
return self._gtTotal
|
||||
|
||||
@property
|
||||
def resTotal(self):
|
||||
return self._resTotal
|
||||
|
||||
def __str__(self):
|
||||
return "true: {}, gt: {}, res: {}".format(self._truePos, self._gtTotal, self._resTotal)
|
||||
@@ -0,0 +1,404 @@
|
||||
"""
|
||||
Evaluation of -.tar.gz file.
|
||||
Yu Fang - March 2019
|
||||
"""
|
||||
|
||||
import os
|
||||
import xml.dom.minidom
|
||||
|
||||
# from eval import eval
|
||||
|
||||
reg_gt_path = os.path.abspath("data/test")
|
||||
reg_gt_path_archival = os.path.abspath("data/test")
|
||||
reg_gt_path_modern = os.path.abspath("data/test")
|
||||
str_gt_path_1 = os.path.abspath("data/test")
|
||||
str_gt_path_2 = os.path.abspath("data/test")
|
||||
str_gt_path_archival = os.path.abspath("data/test")
|
||||
str_gt_path_modern = os.path.abspath("data/test")
|
||||
|
||||
import xml.dom.minidom
|
||||
# from functools import cmp_to_key
|
||||
from os.path import join as osj
|
||||
from .data_structure import *
|
||||
|
||||
|
||||
class eval:
|
||||
STR = "-str"
|
||||
REG = "-reg"
|
||||
DEFAULT_ENCODING = "UTF-8"
|
||||
# reg_gt_path = "./annotations/trackA/"
|
||||
# str_gt_path = "./annotations/trackB/"
|
||||
# reg_gt_path = os.path.abspath("data/test")
|
||||
# reg_gt_path_archival = os.path.abspath("data/test")
|
||||
# reg_gt_path_modern = os.path.abspath("data/test")
|
||||
# str_gt_path_1 = os.path.abspath("data/test")
|
||||
# str_gt_path_2 = os.path.abspath("data/test")
|
||||
# str_gt_path_archival = os.path.abspath("data/test")
|
||||
# str_gt_path_modern = os.path.abspath("data/test")
|
||||
|
||||
# dummyDom = xml.dom.minidom.parse("./dummyXML.xml")
|
||||
|
||||
def __init__(self, track, res_path):
|
||||
self.return_result = None
|
||||
self.reg = True
|
||||
self.str = False
|
||||
|
||||
self.resultFile = res_path
|
||||
self.inPrefix = os.path.split(res_path)[-1].split(".")[0][:-7]
|
||||
|
||||
if track == "-trackA":
|
||||
self.reg = True
|
||||
self.GTFile = osj(reg_gt_path, self.inPrefix + ".xml")
|
||||
# self.GTFile = osj(self.reg_gt_path, self.inPrefix)
|
||||
elif track == "-trackA1": # archival documents
|
||||
self.reg = True
|
||||
self.GTFile = osj(reg_gt_path_archival, self.inPrefix + ".xml")
|
||||
elif track == "-trackA2": # modern documents
|
||||
self.reg = True
|
||||
self.GTFile = osj(reg_gt_path_modern, self.inPrefix + ".xml")
|
||||
elif track == "-trackB1":
|
||||
self.str = True
|
||||
self.GTFile = osj(str_gt_path_1, self.inPrefix + ".xml")
|
||||
# self.GTFile = osj(self.str_gt_path_1, self.inPrefix)
|
||||
elif track == "-trackB2":
|
||||
self.str = True
|
||||
self.GTFile = osj(str_gt_path_2, self.inPrefix + ".xml")
|
||||
# print(self.GTFile)
|
||||
# self.GTFile = osj(self.str_gt_path_2, self.inPrefix)
|
||||
elif track == "-trackB2_a":
|
||||
self.str = True
|
||||
self.GTFile = osj(str_gt_path_archival, self.inPrefix + ".xml")
|
||||
elif track == "-trackB2_m":
|
||||
self.str = True
|
||||
self.GTFile = osj(str_gt_path_modern, self.inPrefix + ".xml")
|
||||
else:
|
||||
print(track)
|
||||
print("Not a valid track, please check your spelling.")
|
||||
|
||||
# self.resultFile = res_path
|
||||
# self.inPrefix = os.path.split(res_path)[-1].split("-")[0]
|
||||
|
||||
# if self.str:
|
||||
# # self.GTFile = osj(self.str_gt_path, self.inPrefix + "-str.xml")
|
||||
# self.GTFile = osj(self.str_gt_path, self.inPrefix + ".xml")
|
||||
# elif self.reg:
|
||||
# # self.GTFile = osj(self.reg_gt_path, self.inPrefix + "-reg.xml")
|
||||
# self.GTFile = osj(self.reg_gt_path, self.inPrefix + ".xml")
|
||||
# else:
|
||||
# print("Not a valid track, please check your spelling.")
|
||||
|
||||
self.gene_ret_lst()
|
||||
|
||||
@property
|
||||
def result(self):
|
||||
return self.return_result
|
||||
|
||||
def gene_ret_lst(self):
|
||||
ret_lst = []
|
||||
for iou in [0.6, 0.7, 0.8, 0.9]:
|
||||
temp = self.compute_retVal(iou)
|
||||
ret_lst.append(temp)
|
||||
# ret_lst.append(self.compute_retVal(iou))
|
||||
|
||||
ret_lst.append(self.inPrefix + ".xml")
|
||||
# ret_lst.append(self.inPrefix)
|
||||
# print("Done processing {}\n".format(self.resultFile))
|
||||
self.return_result = ret_lst
|
||||
|
||||
def compute_retVal(self, iou):
|
||||
gt_dom = xml.dom.minidom.parse(self.GTFile)
|
||||
# incorrect submission format handling
|
||||
try:
|
||||
result_dom = xml.dom.minidom.parse(self.resultFile)
|
||||
except Exception as e:
|
||||
# result_dom = xml.dom.minidom.parse(dummyDom)
|
||||
gt_tables = eval.get_table_list(gt_dom)
|
||||
retVal = ResultStructure(truePos=0, gtTotal=len(gt_tables), resTotal=0)
|
||||
return retVal
|
||||
|
||||
# result_dom = xml.dom.minidom.parse(self.resultFile)
|
||||
if self.reg:
|
||||
ret = self.evaluate_result_reg(gt_dom, result_dom, iou)
|
||||
return ret
|
||||
if self.str:
|
||||
ret = self.evaluate_result_str(gt_dom, result_dom, iou)
|
||||
return ret
|
||||
|
||||
@staticmethod
|
||||
def get_table_list(dom):
|
||||
"""
|
||||
return a list of Table objects corresponding to the table element of the DOM.
|
||||
"""
|
||||
return [Table(_nd) for _nd in dom.documentElement.getElementsByTagName("table")]
|
||||
|
||||
@staticmethod
|
||||
def evaluate_result_reg(gt_dom, result_dom, iou_value):
|
||||
# parse the tables in input elements
|
||||
gt_tables = eval.get_table_list(gt_dom)
|
||||
result_tables = eval.get_table_list(result_dom)
|
||||
# duplicate result table list
|
||||
remaining_tables = result_tables.copy()
|
||||
|
||||
# map the tables in gt and result file
|
||||
table_matches = [] # @param: table_matches - list of mapping of tables in gt and res file, in order (gt, res)
|
||||
for gtt in gt_tables:
|
||||
for rest in remaining_tables:
|
||||
if gtt.compute_table_iou(rest) >= iou_value:
|
||||
remaining_tables.remove(rest)
|
||||
table_matches.append((gtt, rest))
|
||||
break
|
||||
|
||||
assert len(table_matches) <= len(gt_tables)
|
||||
assert len(table_matches) <= len(result_tables)
|
||||
|
||||
retVal = ResultStructure(truePos=len(table_matches), gtTotal=len(gt_tables), resTotal=len(result_tables))
|
||||
return retVal
|
||||
|
||||
@staticmethod
|
||||
def evaluate_result_str(gt_dom, result_dom, iou_value, table_iou_value=0.8):
|
||||
# parse the tables in input elements
|
||||
gt_tables = eval.get_table_list(gt_dom)
|
||||
result_tables = eval.get_table_list(result_dom)
|
||||
|
||||
# duplicate result table list
|
||||
remaining_tables = result_tables.copy()
|
||||
gt_remaining = gt_tables.copy()
|
||||
|
||||
# map the tables in gt and result file
|
||||
table_matches = [] # @param: table_matches - list of mapping of tables in gt and res file, in order (gt, res)
|
||||
for gtt in gt_remaining:
|
||||
for rest in remaining_tables:
|
||||
# note: for structural analysis, use 0.8 for table mapping
|
||||
if gtt.compute_table_iou(rest) >= table_iou_value:
|
||||
table_matches.append((gtt, rest))
|
||||
remaining_tables.remove(rest) # unsafe... should be ok with the break below
|
||||
gt_remaining.remove(gtt)
|
||||
break
|
||||
|
||||
total_gt_relation, total_res_relation, total_correct_relation = 0, 0, 0
|
||||
for gt_table, ress_table in table_matches:
|
||||
|
||||
# set up the cell mapping for matching tables
|
||||
cell_mapping = gt_table.find_cell_mapping(ress_table, iou_value)
|
||||
# set up the adj relations, convert the one for result table to a dictionary for faster searching
|
||||
gt_AR = gt_table.find_adj_relations()
|
||||
total_gt_relation += len(gt_AR)
|
||||
|
||||
res_AR = ress_table.find_adj_relations()
|
||||
total_res_relation += len(res_AR)
|
||||
|
||||
if False: # for DEBUG
|
||||
Table.printCellMapping(cell_mapping)
|
||||
Table.printAdjacencyRelationList(gt_AR, "GT")
|
||||
Table.printAdjacencyRelationList(res_AR, "run")
|
||||
|
||||
# Now map GT adjacency relations to result
|
||||
lMappedAR = []
|
||||
for ar in gt_AR:
|
||||
try:
|
||||
resFromCell = cell_mapping[ar.fromText]
|
||||
resToCell = cell_mapping[ar.toText]
|
||||
# make a mapped adjacency relation
|
||||
lMappedAR.append(AdjRelation(resFromCell, resToCell, ar.direction))
|
||||
except:
|
||||
# no mapping is possible
|
||||
pass
|
||||
|
||||
# compare two list of adjacency relation
|
||||
correct_dect = 0
|
||||
for ar1 in res_AR:
|
||||
for ar2 in lMappedAR:
|
||||
if ar1.isEqual(ar2):
|
||||
correct_dect += 1
|
||||
break
|
||||
|
||||
total_correct_relation += correct_dect
|
||||
|
||||
# handle gt_relations in unmatched gt table
|
||||
for gtt_remain in gt_remaining:
|
||||
total_gt_relation += len(gtt_remain.find_adj_relations())
|
||||
|
||||
# handle gt_relation in unmatched res table
|
||||
for res_remain in remaining_tables:
|
||||
total_res_relation += len(res_remain.find_adj_relations())
|
||||
|
||||
retVal = ResultStructure(truePos=total_correct_relation, gtTotal=total_gt_relation, resTotal=total_res_relation)
|
||||
return retVal
|
||||
|
||||
# calculate the gt adj_relations of the missing file
|
||||
# @param: file_lst - list of missing ground truth file
|
||||
# @param: cur_gt_num - current total of ground truth objects (tables / cells)
|
||||
def process_missing_files(track, gt_file_lst, cur_gt_num):
|
||||
if track in ["-trackA", "-trackA1", "-trackA2"]:
|
||||
gt_file_lst_full = [osj(reg_gt_path, filename) for filename in gt_file_lst]
|
||||
for file in gt_file_lst_full:
|
||||
if os.path.split(file)[-1].split(".")[-1] == "xml":
|
||||
gt_dom = xml.dom.minidom.parse(file)
|
||||
gt_root = gt_dom.documentElement
|
||||
# tables = []
|
||||
table_elements = gt_root.getElementsByTagName("table")
|
||||
for res_table in table_elements:
|
||||
# t = Table(res_table)
|
||||
# tables.append(t)
|
||||
cur_gt_num += 1
|
||||
return cur_gt_num
|
||||
elif track == "-trackB1":
|
||||
gt_file_lst_full = [osj(str_gt_path_1, filename) for filename in gt_file_lst]
|
||||
for file in gt_file_lst_full:
|
||||
if os.path.split(file)[-1].split(".")[-1] == "xml":
|
||||
gt_dom = xml.dom.minidom.parse(file)
|
||||
gt_root = gt_dom.documentElement
|
||||
tables = []
|
||||
table_elements = gt_root.getElementsByTagName("table")
|
||||
for res_table in table_elements:
|
||||
t = Table(res_table)
|
||||
tables.append(t)
|
||||
for table in tables:
|
||||
cur_gt_num += len(table.find_adj_relations())
|
||||
return cur_gt_num
|
||||
elif track == "-trackB2":
|
||||
gt_file_lst_full = [osj(str_gt_path_2, filename) for filename in gt_file_lst]
|
||||
for file in gt_file_lst_full:
|
||||
if os.path.split(file)[-1].split(".")[-1] == "xml":
|
||||
gt_dom = xml.dom.minidom.parse(file)
|
||||
gt_root = gt_dom.documentElement
|
||||
tables = []
|
||||
table_elements = gt_root.getElementsByTagName("table")
|
||||
for res_table in table_elements:
|
||||
t = Table(res_table)
|
||||
tables.append(t)
|
||||
for table in tables:
|
||||
cur_gt_num += len(table.find_adj_relations())
|
||||
return cur_gt_num
|
||||
|
||||
def calc(F1):
|
||||
sum_a = 0.6 * F1[0] + 0.7 * F1[1] + 0.8 * F1[2] + 0.9 * F1[3]
|
||||
sum_b = 0.6 + 0.7 + 0.8 + 0.9
|
||||
|
||||
return sum_a / sum_b
|
||||
|
||||
def calc_table_score(result_path):
|
||||
# measure = eval(*sys.argv[1:])
|
||||
|
||||
gt_file_lst = os.listdir(reg_gt_path_archival)
|
||||
track = "-trackA1"
|
||||
untar_path = result_path
|
||||
|
||||
res_lst = []
|
||||
for root, files, dirs in os.walk(untar_path):
|
||||
for name in dirs:
|
||||
if name.split(".")[-1] == "xml":
|
||||
cur_filepath = osj(os.path.abspath(root), name)
|
||||
res_lst.append(eval(track, cur_filepath))
|
||||
# printing for debug
|
||||
# print("Processing... {}".format(name))
|
||||
# print("DONE WITH FILE PROCESSING\n")
|
||||
# note: results are stored as list of each when iou at [0.6, 0.7, 0.8, 0.9, gt_filename]
|
||||
# gt number should be the same for all files
|
||||
gt_num = 0
|
||||
correct_six, res_six = 0, 0
|
||||
correct_seven, res_seven = 0, 0
|
||||
correct_eight, res_eight = 0, 0
|
||||
correct_nine, res_nine = 0, 0
|
||||
|
||||
|
||||
for each_file in res_lst:
|
||||
# print(each_file)
|
||||
try:
|
||||
gt_file_lst.remove(each_file.result[-1])
|
||||
if each_file.result[-1].replace('.xml', '.jpg') in gt_file_lst:
|
||||
gt_file_lst.remove(each_file.result[-1].replace('.xml', '.jpg'))
|
||||
correct_six += each_file.result[0].truePos
|
||||
gt_num += each_file.result[0].gtTotal
|
||||
res_six += each_file.result[0].resTotal
|
||||
# print("{} {} {}".format(each_file.result[0].truePos, each_file.result[0].gtTotal, each_file.result[0].resTotal))
|
||||
|
||||
correct_seven += each_file.result[1].truePos
|
||||
res_seven += each_file.result[1].resTotal
|
||||
|
||||
correct_eight += each_file.result[2].truePos
|
||||
res_eight += each_file.result[2].resTotal
|
||||
|
||||
correct_nine += each_file.result[3].truePos
|
||||
res_nine += each_file.result[3].resTotal
|
||||
except:
|
||||
print("Error occur in processing result list.")
|
||||
print(each_file.result[-1])
|
||||
break
|
||||
# print(each_file.result[-1])
|
||||
# print(each_file)
|
||||
|
||||
# for file in gt_file_lst:
|
||||
# if file.split(".") != "xml":
|
||||
# gt_file_lst.remove(file)
|
||||
# # print(gt_file_lst)
|
||||
|
||||
for i in range(len(gt_file_lst) - 1, -1, -1):
|
||||
if gt_file_lst[i].split(".")[-1] != "xml":
|
||||
del gt_file_lst[i]
|
||||
|
||||
if len(gt_file_lst) > 0:
|
||||
print("\nWarning: missing result annotations for file: {}\n".format(gt_file_lst))
|
||||
gt_total = process_missing_files(track, gt_file_lst, gt_num)
|
||||
else:
|
||||
gt_total = gt_num
|
||||
|
||||
|
||||
try:
|
||||
# print("Evaluation of {}".format(track.replace("-", "")))
|
||||
# iou @ 0.6
|
||||
p_six = correct_six / res_six
|
||||
r_six = correct_six / gt_total
|
||||
f1_six = 2 * p_six * r_six / (p_six + r_six)
|
||||
print("IOU @ 0.6 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_six, r_six, f1_six))
|
||||
print("correct: {}, gt: {}, res: {}\n".format(correct_six, gt_total, res_six))
|
||||
|
||||
# iou @ 0.7
|
||||
p_seven = correct_seven / res_seven
|
||||
r_seven = correct_seven / gt_total
|
||||
f1_seven = 2 * p_seven * r_seven / (p_seven + r_seven)
|
||||
print("IOU @ 0.7 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_seven, r_seven, f1_seven))
|
||||
print("correct: {}, gt: {}, res: {}\n".format(correct_seven, gt_total, res_seven))
|
||||
|
||||
# iou @ 0.8
|
||||
p_eight = correct_eight / res_eight
|
||||
r_eight = correct_eight / gt_total
|
||||
f1_eight = 2 * p_eight * r_eight / (p_eight + r_eight)
|
||||
print("IOU @ 0.8 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_eight, r_eight, f1_eight))
|
||||
print("correct: {}, gt: {}, res: {}\n".format(correct_eight, gt_total, res_eight))
|
||||
|
||||
# iou @ 0.9
|
||||
p_nine = correct_nine / res_nine
|
||||
r_nine = correct_nine / gt_total
|
||||
f1_nine = 2 * p_nine * r_nine / (p_nine + r_nine)
|
||||
print("IOU @ 0.9 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_nine, r_nine, f1_nine))
|
||||
print("correct: {}, gt: {}, res: {}".format(correct_nine, gt_total, res_nine))
|
||||
|
||||
F1 = [f1_six, f1_seven, f1_eight, f1_nine]
|
||||
wF1 = calc(F1)
|
||||
|
||||
print("Average weight F1: {}".format(wF1))
|
||||
|
||||
return {
|
||||
'p_six':p_six * 100,
|
||||
"r_six":r_six * 100,
|
||||
"f1_six":f1_six * 100,
|
||||
"p_seven":p_seven * 100,
|
||||
"r_seven":r_seven * 100,
|
||||
"f1_seven":f1_seven * 100,
|
||||
"p_eight":p_eight * 100,
|
||||
"r_eight":r_eight * 100,
|
||||
"f1_eight":f1_eight * 100,
|
||||
"p_nine":p_nine * 100,
|
||||
"r_nine":r_nine * 100,
|
||||
"f1_nine":f1_nine * 100,
|
||||
"wF1":wF1 * 100
|
||||
}
|
||||
except ZeroDivisionError:
|
||||
print(
|
||||
"Error: zero devision error found, (possible that no adjacency relations are found), please check the file input.")
|
||||
return {"wF1": 0}
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
pass
|
||||
@@ -0,0 +1,69 @@
|
||||
MODEL:
|
||||
MASK_ON: True
|
||||
META_ARCHITECTURE: "GeneralizedRCNN"
|
||||
PIXEL_MEAN: [127.5, 127.5, 127.5]
|
||||
PIXEL_STD: [127.5, 127.5, 127.5]
|
||||
BACKBONE:
|
||||
NAME: "build_vit_fpn_backbone"
|
||||
VIT:
|
||||
OUT_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
DROP_PATH: 0.1
|
||||
IMG_SIZE: [224,224]
|
||||
POS_TYPE: "abs"
|
||||
FPN:
|
||||
IN_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
ANCHOR_GENERATOR:
|
||||
SIZES: [[32], [64], [128], [256], [512]] # One size for each in feature map
|
||||
ASPECT_RATIOS: [[0.5, 1.0, 2.0]] # Three aspect ratios (same for all in feature maps)
|
||||
RPN:
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
|
||||
PRE_NMS_TOPK_TRAIN: 2000 # Per FPN level
|
||||
PRE_NMS_TOPK_TEST: 1000 # Per FPN level
|
||||
# Detectron1 uses 2000 proposals per-batch,
|
||||
# (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
|
||||
# which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
|
||||
POST_NMS_TOPK_TRAIN: 1000
|
||||
POST_NMS_TOPK_TEST: 1000
|
||||
ROI_HEADS:
|
||||
NAME: "StandardROIHeads"
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5"]
|
||||
NUM_CLASSES: 1
|
||||
ROI_BOX_HEAD:
|
||||
NAME: "FastRCNNConvFCHead"
|
||||
NUM_FC: 2
|
||||
POOLER_RESOLUTION: 7
|
||||
ROI_MASK_HEAD:
|
||||
NAME: "MaskRCNNConvUpsampleHead"
|
||||
NUM_CONV: 4
|
||||
POOLER_RESOLUTION: 14
|
||||
DATASETS:
|
||||
TRAIN: ("icdar2019_train",)
|
||||
TEST: ("icdar2019_test",)
|
||||
SOLVER:
|
||||
LR_SCHEDULER_NAME: "WarmupCosineLR"
|
||||
AMP:
|
||||
ENABLED: True
|
||||
OPTIMIZER: "ADAMW"
|
||||
BACKBONE_MULTIPLIER: 1.0
|
||||
CLIP_GRADIENTS:
|
||||
ENABLED: True
|
||||
CLIP_TYPE: "full_model"
|
||||
CLIP_VALUE: 1.0
|
||||
NORM_TYPE: 2.0
|
||||
WARMUP_FACTOR: 0.01
|
||||
BASE_LR: 0.0002
|
||||
WEIGHT_DECAY: 0.05
|
||||
IMS_PER_BATCH: 32
|
||||
INPUT:
|
||||
CROP:
|
||||
ENABLED: True
|
||||
TYPE: "absolute_range"
|
||||
SIZE: (384, 600)
|
||||
MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
|
||||
FORMAT: "RGB"
|
||||
DATALOADER:
|
||||
FILTER_EMPTY_ANNOTATIONS: False
|
||||
VERSION: 2
|
||||
AUG:
|
||||
DETR: True
|
||||
SEED: 42
|
||||
@@ -0,0 +1,21 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
|
||||
VIT:
|
||||
NAME: "dit_base_patch16"
|
||||
ROI_HEADS:
|
||||
NAME: CascadeROIHeads
|
||||
ROI_BOX_HEAD:
|
||||
CLS_AGNOSTIC_BBOX_REG: True
|
||||
RPN:
|
||||
POST_NMS_TOPK_TRAIN: 2000
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.00005
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
@@ -0,0 +1,25 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
|
||||
VIT:
|
||||
NAME: "dit_large_patch16"
|
||||
OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
DROP_PATH: 0.2
|
||||
FPN:
|
||||
IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
ROI_HEADS:
|
||||
NAME: CascadeROIHeads
|
||||
ROI_BOX_HEAD:
|
||||
CLS_AGNOSTIC_BBOX_REG: True
|
||||
RPN:
|
||||
POST_NMS_TOPK_TRAIN: 2000
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.00005
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
@@ -0,0 +1,15 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
|
||||
VIT:
|
||||
NAME: "dit_base_patch16"
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.00005
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
@@ -0,0 +1,19 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
|
||||
VIT:
|
||||
NAME: "dit_large_patch16"
|
||||
OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
DROP_PATH: 0.2
|
||||
FPN:
|
||||
IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.00005
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
@@ -0,0 +1,80 @@
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
|
||||
from ditod import add_vit_config
|
||||
|
||||
import torch
|
||||
|
||||
from detectron2.config import get_cfg
|
||||
from detectron2.utils.visualizer import ColorMode, Visualizer
|
||||
from detectron2.data import MetadataCatalog
|
||||
from detectron2.engine import DefaultPredictor
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Detectron2 inference script")
|
||||
parser.add_argument(
|
||||
"--image_path",
|
||||
help="Path to input image",
|
||||
type=str,
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_file_name",
|
||||
help="Name of the output visualization file.",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--config-file",
|
||||
default="configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml",
|
||||
metavar="FILE",
|
||||
help="path to config file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--opts",
|
||||
help="Modify config options using the command-line 'KEY VALUE' pairs",
|
||||
default=[],
|
||||
nargs=argparse.REMAINDER,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Step 1: instantiate config
|
||||
cfg = get_cfg()
|
||||
add_vit_config(cfg)
|
||||
cfg.merge_from_file(args.config_file)
|
||||
|
||||
# Step 2: add model weights URL to config
|
||||
cfg.merge_from_list(args.opts)
|
||||
|
||||
# Step 3: set device
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
cfg.MODEL.DEVICE = device
|
||||
|
||||
# Step 4: define model
|
||||
predictor = DefaultPredictor(cfg)
|
||||
|
||||
# Step 5: run inference
|
||||
img = cv2.imread(args.image_path)
|
||||
|
||||
md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
|
||||
if cfg.DATASETS.TEST[0]=='icdar2019_test':
|
||||
md.set(thing_classes=["table"])
|
||||
else:
|
||||
md.set(thing_classes=["text","title","list","table","figure"])
|
||||
|
||||
output = predictor(img)["instances"]
|
||||
v = Visualizer(img[:, :, ::-1],
|
||||
md,
|
||||
scale=1.0,
|
||||
instance_mode=ColorMode.SEGMENTATION)
|
||||
result = v.draw_instance_predictions(output.to("cpu"))
|
||||
result_image = result.get_image()[:, :, ::-1]
|
||||
|
||||
# step 6: save
|
||||
cv2.imwrite(args.output_file_name, result_image)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
MODEL:
|
||||
MASK_ON: True
|
||||
META_ARCHITECTURE: "GeneralizedRCNN"
|
||||
PIXEL_MEAN: [123.675, 116.280, 103.530]
|
||||
PIXEL_STD: [58.395, 57.120, 57.375]
|
||||
BACKBONE:
|
||||
NAME: "build_vit_fpn_backbone"
|
||||
VIT:
|
||||
OUT_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
DROP_PATH: 0.1
|
||||
IMG_SIZE: [224,224]
|
||||
POS_TYPE: "abs"
|
||||
FPN:
|
||||
IN_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
ANCHOR_GENERATOR:
|
||||
SIZES: [[32], [64], [128], [256], [512]] # One size for each in feature map
|
||||
ASPECT_RATIOS: [[0.5, 1.0, 2.0]] # Three aspect ratios (same for all in feature maps)
|
||||
RPN:
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
|
||||
PRE_NMS_TOPK_TRAIN: 2000 # Per FPN level
|
||||
PRE_NMS_TOPK_TEST: 1000 # Per FPN level
|
||||
# Detectron1 uses 2000 proposals per-batch,
|
||||
# (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
|
||||
# which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
|
||||
POST_NMS_TOPK_TRAIN: 1000
|
||||
POST_NMS_TOPK_TEST: 1000
|
||||
ROI_HEADS:
|
||||
NAME: "StandardROIHeads"
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5"]
|
||||
NUM_CLASSES: 5
|
||||
ROI_BOX_HEAD:
|
||||
NAME: "FastRCNNConvFCHead"
|
||||
NUM_FC: 2
|
||||
POOLER_RESOLUTION: 7
|
||||
ROI_MASK_HEAD:
|
||||
NAME: "MaskRCNNConvUpsampleHead"
|
||||
NUM_CONV: 4
|
||||
POOLER_RESOLUTION: 14
|
||||
DATASETS:
|
||||
TRAIN: ("publaynet_train",)
|
||||
TEST: ("publaynet_val",)
|
||||
SOLVER:
|
||||
LR_SCHEDULER_NAME: "WarmupCosineLR"
|
||||
AMP:
|
||||
ENABLED: True
|
||||
OPTIMIZER: "ADAMW"
|
||||
BACKBONE_MULTIPLIER: 1.0
|
||||
CLIP_GRADIENTS:
|
||||
ENABLED: True
|
||||
CLIP_TYPE: "full_model"
|
||||
CLIP_VALUE: 1.0
|
||||
NORM_TYPE: 2.0
|
||||
WARMUP_FACTOR: 0.01
|
||||
BASE_LR: 0.0004
|
||||
WEIGHT_DECAY: 0.05
|
||||
IMS_PER_BATCH: 32
|
||||
INPUT:
|
||||
CROP:
|
||||
ENABLED: True
|
||||
TYPE: "absolute_range"
|
||||
SIZE: (384, 600)
|
||||
MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
|
||||
FORMAT: "RGB"
|
||||
DATALOADER:
|
||||
FILTER_EMPTY_ANNOTATIONS: False
|
||||
VERSION: 2
|
||||
AUG:
|
||||
DETR: True
|
||||
SEED: 42
|
||||
@@ -0,0 +1,20 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
|
||||
VIT:
|
||||
NAME: "dit_base_patch16"
|
||||
ROI_HEADS:
|
||||
NAME: CascadeROIHeads
|
||||
ROI_BOX_HEAD:
|
||||
CLS_AGNOSTIC_BBOX_REG: True
|
||||
RPN:
|
||||
POST_NMS_TOPK_TRAIN: 2000
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 2000
|
||||
TEST:
|
||||
EVAL_PERIOD: 2000
|
||||
@@ -0,0 +1,28 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
|
||||
VIT:
|
||||
NAME: "dit_large_patch16"
|
||||
OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
DROP_PATH: 0.2
|
||||
FPN:
|
||||
IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
ROI_HEADS:
|
||||
NAME: CascadeROIHeads
|
||||
ROI_BOX_HEAD:
|
||||
CLS_AGNOSTIC_BBOX_REG: True
|
||||
RPN:
|
||||
POST_NMS_TOPK_TRAIN: 2000
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 2000
|
||||
BASE_LR: 0.0001
|
||||
STEPS: (40000, 53333)
|
||||
AMP:
|
||||
ENABLED: False
|
||||
TEST:
|
||||
EVAL_PERIOD: 2000
|
||||
@@ -0,0 +1,15 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
|
||||
VIT:
|
||||
NAME: "dit_base_patch16"
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 2000
|
||||
TEST:
|
||||
EVAL_PERIOD: 2000
|
||||
OUTPUT_DIR: $AMLT_OUTPUT_DIR
|
||||
@@ -0,0 +1,22 @@
|
||||
_BASE_: "../Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
|
||||
VIT:
|
||||
NAME: "dit_large_patch16"
|
||||
OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
DROP_PATH: 0.2
|
||||
FPN:
|
||||
IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 2000
|
||||
BASE_LR: 0.0001
|
||||
AMP:
|
||||
ENABLED: False
|
||||
TEST:
|
||||
EVAL_PERIOD: 2000
|
||||
OUTPUT_DIR: "output/publaynet/mask_rcnn/dit_base_multistep_3x_ms"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 319 KiB |
@@ -0,0 +1,123 @@
|
||||
#!/usr/bin/env python
|
||||
# --------------------------------------------------------------------------------
|
||||
# MPViT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Detection Training Script for MPViT.
|
||||
"""
|
||||
|
||||
import os
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
|
||||
from typing import Any, Dict, List, Set
|
||||
|
||||
from detectron2.data import build_detection_train_loader
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
from detectron2.config import get_cfg
|
||||
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch
|
||||
from detectron2.evaluation import COCOEvaluator
|
||||
from detectron2.solver.build import maybe_add_gradient_clipping
|
||||
|
||||
from ditod import add_vit_config
|
||||
from ditod import DetrDatasetMapper
|
||||
|
||||
from detectron2.data.datasets import register_coco_instances
|
||||
import logging
|
||||
from detectron2.utils.logger import setup_logger
|
||||
from detectron2.utils import comm
|
||||
from detectron2.engine.defaults import create_ddp_model
|
||||
import weakref
|
||||
from detectron2.engine.train_loop import AMPTrainer, SimpleTrainer
|
||||
from ditod import MyDetectionCheckpointer, ICDAREvaluator
|
||||
from ditod import MyTrainer
|
||||
|
||||
|
||||
def setup(args):
|
||||
"""
|
||||
Create configs and perform basic setups.
|
||||
"""
|
||||
cfg = get_cfg()
|
||||
# add_coat_config(cfg)
|
||||
add_vit_config(cfg)
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
default_setup(cfg, args)
|
||||
return cfg
|
||||
|
||||
|
||||
def main(args):
|
||||
"""
|
||||
register publaynet first
|
||||
"""
|
||||
register_coco_instances(
|
||||
"publaynet_train",
|
||||
{},
|
||||
"./publaynet_data/train.json",
|
||||
"./publaynet_data/train"
|
||||
)
|
||||
|
||||
register_coco_instances(
|
||||
"publaynet_val",
|
||||
{},
|
||||
"./publaynet_data/val.json",
|
||||
"./publaynet_data/val"
|
||||
)
|
||||
|
||||
register_coco_instances(
|
||||
"icdar2019_train",
|
||||
{},
|
||||
"data/train.json",
|
||||
"data/train"
|
||||
)
|
||||
|
||||
register_coco_instances(
|
||||
"icdar2019_test",
|
||||
{},
|
||||
"data/test.json",
|
||||
"data/test"
|
||||
)
|
||||
|
||||
cfg = setup(args)
|
||||
|
||||
if args.eval_only:
|
||||
model = MyTrainer.build_model(cfg)
|
||||
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
|
||||
cfg.MODEL.WEIGHTS, resume=args.resume
|
||||
)
|
||||
res = MyTrainer.test(cfg, model)
|
||||
return res
|
||||
|
||||
trainer = MyTrainer(cfg)
|
||||
trainer.resume_or_load(resume=args.resume)
|
||||
return trainer.train()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = default_argument_parser()
|
||||
parser.add_argument("--debug", action="store_true", help="enable debug mode")
|
||||
args = parser.parse_args()
|
||||
print("Command Line Args:", args)
|
||||
|
||||
if args.debug:
|
||||
import debugpy
|
||||
|
||||
print("Enabling attach starts.")
|
||||
debugpy.listen(address=('0.0.0.0', 9310))
|
||||
debugpy.wait_for_client()
|
||||
print("Enabling attach ends.")
|
||||
|
||||
launch(
|
||||
main,
|
||||
args.num_gpus,
|
||||
num_machines=args.num_machines,
|
||||
machine_rank=args.machine_rank,
|
||||
dist_url=args.dist_url,
|
||||
args=(args,),
|
||||
)
|
||||
@@ -0,0 +1,14 @@
|
||||
torch==1.9.0
|
||||
torchvision==0.10.0
|
||||
timm==0.5.4
|
||||
Pillow
|
||||
blobfile
|
||||
mypy
|
||||
numpy
|
||||
pytest
|
||||
requests
|
||||
einops
|
||||
tensorboardX
|
||||
deepspeed==0.4.0
|
||||
scipy
|
||||
opencv-python
|
||||
@@ -0,0 +1,62 @@
|
||||
# DiT for Text Detection
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/45008728/163219997-90d15c1b-e1d1-4bb3-ae46-774e54b89dc6.png" width="500" /><img src="https://user-images.githubusercontent.com/45008728/163220437-ab6a3fd2-0a4f-49c5-810c-e05dda7eb9e1.png" width="500"/> Model outputs with FUNSD
|
||||
</div>
|
||||
|
||||
## Fine-tuned models on FUNSD
|
||||
We summarize the validation results as follows. We also provide the fine-tuned weights.
|
||||
|
||||
| name | initialized checkpoint | detection algorithm | F1 | weight |
|
||||
|------------|:----------------------------------------|:----------:|-------------------|-----|
|
||||
| DiT-base-syn | [dit_base_patch16_224_syn](https://layoutlm.blob.core.windows.net/dit/dit-fts/td-syn_dit-b_mrcnn.pth) | Mask R-CNN | 94.25 | [link](https://layoutlm.blob.core.windows.net/dit/dit-fts/funsd_dit-b_mrcnn.pth) |
|
||||
| DiT-large-syn | [dit_large_patch16_224_syn](https://layoutlm.blob.core.windows.net/dit/dit-fts/td-syn_dit-l_mrcnn.pth) | Mask R-CNN | 94.29 | [link](https://layoutlm.blob.core.windows.net/dit/dit-fts/funsd_dit-l_mrcnn.pth) |
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
### Data Preparation
|
||||
|
||||
Follow [these steps](https://mmocr.readthedocs.io/en/latest/datasets/det.html#funsd) to download and process the FUNSD. The resulting directory structure looks like the following:
|
||||
```
|
||||
│── data
|
||||
│ ├── annotations
|
||||
│ ├── imgs
|
||||
│ ├── instances_test.json
|
||||
│ └── instances_training.json
|
||||
|
||||
```
|
||||
### Training
|
||||
The following command provide example to train the Mask R-CNN with DiT backbone on 8 32GB Nvidia V100 GPUs.
|
||||
|
||||
The config files can be found in `configs`.
|
||||
|
||||
```bash
|
||||
python train_net.py --config-file configs/mask_rcnn_dit_base.yaml --num-gpus 8 --resume MODEL.WEIGHTS path/to/model OUTPUT_DIR path/to/output
|
||||
```
|
||||
|
||||
### Evaluation
|
||||
The following commands provide examples to evaluate the fine-tuned checkpoint of DiT-Base with Mask R-CNN.
|
||||
|
||||
```bash
|
||||
python train_net.py --config-file configs/mask_rcnn_dit_base.yaml --eval-only --num-gpus 8 --resume MODEL.WEIGHTS path/to/model OUTPUT_DIR path/to/output
|
||||
```
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider citing our work:
|
||||
```
|
||||
@misc{li2022dit,
|
||||
title={DiT: Self-supervised Pre-training for Document Image Transformer},
|
||||
author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
|
||||
year={2022},
|
||||
eprint={2203.02378},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Acknowledgment
|
||||
Thanks to [Detectron2](https://github.com/facebookresearch/detectron2) for Mask R-CNN implementation and [MMOCR](https://github.com/open-mmlab/mmocr) for the data preprocessing implementation of the FUNSD
|
||||
@@ -0,0 +1,69 @@
|
||||
MODEL:
|
||||
MASK_ON: True
|
||||
META_ARCHITECTURE: "GeneralizedRCNN"
|
||||
PIXEL_MEAN: [127.5, 127.5, 127.5]
|
||||
PIXEL_STD: [127.5, 127.5, 127.5]
|
||||
BACKBONE:
|
||||
NAME: "build_vit_fpn_backbone"
|
||||
VIT:
|
||||
OUT_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
DROP_PATH: 0.1
|
||||
IMG_SIZE: [224,224]
|
||||
POS_TYPE: "abs"
|
||||
FPN:
|
||||
IN_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
|
||||
ANCHOR_GENERATOR:
|
||||
SIZES: [[32], [64], [128], [256], [512]] # One size for each in feature map
|
||||
ASPECT_RATIOS: [[0.5, 1.0, 2.0]] # Three aspect ratios (same for all in feature maps)
|
||||
RPN:
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
|
||||
PRE_NMS_TOPK_TRAIN: 2000 # Per FPN level
|
||||
PRE_NMS_TOPK_TEST: 1000 # Per FPN level
|
||||
# Detectron1 uses 2000 proposals per-batch,
|
||||
# (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
|
||||
# which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
|
||||
POST_NMS_TOPK_TRAIN: 1000
|
||||
POST_NMS_TOPK_TEST: 1000
|
||||
ROI_HEADS:
|
||||
NAME: "StandardROIHeads"
|
||||
IN_FEATURES: ["p2", "p3", "p4", "p5"]
|
||||
NUM_CLASSES: 1
|
||||
ROI_BOX_HEAD:
|
||||
NAME: "FastRCNNConvFCHead"
|
||||
NUM_FC: 2
|
||||
POOLER_RESOLUTION: 7
|
||||
ROI_MASK_HEAD:
|
||||
NAME: "MaskRCNNConvUpsampleHead"
|
||||
NUM_CONV: 4
|
||||
POOLER_RESOLUTION: 14
|
||||
DATASETS:
|
||||
TRAIN: ("funsd_train",)
|
||||
TEST: ("funsd_test",)
|
||||
SOLVER:
|
||||
LR_SCHEDULER_NAME: "WarmupCosineLR"
|
||||
AMP:
|
||||
ENABLED: True
|
||||
OPTIMIZER: "ADAMW"
|
||||
BACKBONE_MULTIPLIER: 1.0
|
||||
CLIP_GRADIENTS:
|
||||
ENABLED: True
|
||||
CLIP_TYPE: "full_model"
|
||||
CLIP_VALUE: 1.0
|
||||
NORM_TYPE: 2.0
|
||||
WARMUP_FACTOR: 0.01
|
||||
BASE_LR: 0.0002
|
||||
WEIGHT_DECAY: 0.05
|
||||
IMS_PER_BATCH: 32
|
||||
INPUT:
|
||||
CROP:
|
||||
ENABLED: True
|
||||
TYPE: "absolute_range"
|
||||
SIZE: (384, 600)
|
||||
MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
|
||||
FORMAT: "RGB"
|
||||
DATALOADER:
|
||||
FILTER_EMPTY_ANNOTATIONS: False
|
||||
VERSION: 2
|
||||
AUG:
|
||||
DETR: True
|
||||
SEED: 42
|
||||
@@ -0,0 +1,20 @@
|
||||
_BASE_: "./Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "path/to/model"
|
||||
VIT:
|
||||
NAME: "dit_base_patch16"
|
||||
ANCHOR_GENERATOR:
|
||||
SIZES: [[4], [8], [16], [32], [64]]
|
||||
ASPECT_RATIOS: [[1.5, 3.5, 6.5]]
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 16
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.0001
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
DETECTIONS_PER_IMAGE: 1000
|
||||
OUTPUT_DIR: $$AMLT_OUTPUT_DIR
|
||||
@@ -0,0 +1,26 @@
|
||||
_BASE_: "./Base-RCNN-FPN.yaml"
|
||||
MODEL:
|
||||
PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
|
||||
PIXEL_STD: [ 127.5, 127.5, 127.5 ]
|
||||
WEIGHTS: "path/to/model"
|
||||
VIT:
|
||||
NAME: "dit_large_patch16"
|
||||
OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
DROP_PATH: 0.2
|
||||
FPN:
|
||||
IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
|
||||
ANCHOR_GENERATOR:
|
||||
SIZES: [[4], [8], [16], [32], [64]]
|
||||
ASPECT_RATIOS: [[1.5, 3.5, 6.5]]
|
||||
SOLVER:
|
||||
WARMUP_ITERS: 1000
|
||||
IMS_PER_BATCH: 8
|
||||
MAX_ITER: 60000
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
BASE_LR: 0.00005
|
||||
AMP:
|
||||
ENABLED: True
|
||||
TEST:
|
||||
EVAL_PERIOD: 1000
|
||||
DETECTIONS_PER_IMAGE: 2000
|
||||
OUTPUT_DIR: $AMLT_OUTPUT_DIR
|
||||
@@ -0,0 +1,15 @@
|
||||
# --------------------------------------------------------------------------------
|
||||
# MPViT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
from .config import add_vit_config
|
||||
from .backbone import build_vit_fpn_backbone
|
||||
from .dataset_mapper import DetrDatasetMapper
|
||||
from .mycheckpointer import MyDetectionCheckpointer
|
||||
from .funsd_evaluation import FUNSDEvaluator
|
||||
from .mytrainer import MyTrainer
|
||||
@@ -0,0 +1,156 @@
|
||||
# --------------------------------------------------------------------------------
|
||||
# VIT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# --------------------------------------------------------------------------------
|
||||
# References:
|
||||
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
||||
# CoaT: https://github.com/mlpc-ucsd/CoaT
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
from detectron2.layers import (
|
||||
ShapeSpec,
|
||||
)
|
||||
from detectron2.modeling import Backbone, BACKBONE_REGISTRY, FPN
|
||||
from detectron2.modeling.backbone.fpn import LastLevelP6P7, LastLevelMaxPool
|
||||
|
||||
from .beit import beit_base_patch16, dit_base_patch16, dit_large_patch16, beit_large_patch16
|
||||
from .deit import deit_base_patch16, mae_base_patch16
|
||||
|
||||
__all__ = [
|
||||
"build_vit_fpn_backbone",
|
||||
]
|
||||
|
||||
|
||||
class VIT_Backbone(Backbone):
|
||||
"""
|
||||
Implement VIT backbone.
|
||||
"""
|
||||
|
||||
def __init__(self, name, out_features, drop_path, img_size, pos_type, model_kwargs):
|
||||
super().__init__()
|
||||
self._out_features = out_features
|
||||
if 'base' in name:
|
||||
self._out_feature_strides = {"layer3": 4, "layer5": 8, "layer7": 16, "layer11": 32}
|
||||
else:
|
||||
self._out_feature_strides = {"layer7": 4, "layer11": 8, "layer15": 16, "layer23": 32}
|
||||
|
||||
if name == 'beit_base_patch16':
|
||||
model_func = beit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == 'dit_base_patch16':
|
||||
model_func = dit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "deit_base_patch16":
|
||||
model_func = deit_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "mae_base_patch16":
|
||||
model_func = mae_base_patch16
|
||||
self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
|
||||
elif name == "dit_large_patch16":
|
||||
model_func = dit_large_patch16
|
||||
self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
|
||||
elif name == "beit_large_patch16":
|
||||
model_func = beit_large_patch16
|
||||
self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
|
||||
else:
|
||||
raise ValueError("Unsupported VIT name yet.")
|
||||
|
||||
if 'beit' in name or 'dit' in name:
|
||||
if pos_type == "abs":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_abs_pos_emb=True,
|
||||
**model_kwargs)
|
||||
elif pos_type == "shared_rel":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_shared_rel_pos_bias=True,
|
||||
**model_kwargs)
|
||||
elif pos_type == "rel":
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
use_rel_pos_bias=True,
|
||||
**model_kwargs)
|
||||
else:
|
||||
raise ValueError()
|
||||
else:
|
||||
self.backbone = model_func(img_size=img_size,
|
||||
out_features=out_features,
|
||||
drop_path_rate=drop_path,
|
||||
**model_kwargs)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
|
||||
|
||||
Returns:
|
||||
dict[str->Tensor]: names and the corresponding features
|
||||
"""
|
||||
assert x.dim() == 4, f"VIT takes an input of shape (N, C, H, W). Got {x.shape} instead!"
|
||||
return self.backbone.forward_features(x)
|
||||
|
||||
def output_shape(self):
|
||||
return {
|
||||
name: ShapeSpec(
|
||||
channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
|
||||
)
|
||||
for name in self._out_features
|
||||
}
|
||||
|
||||
|
||||
def build_VIT_backbone(cfg):
|
||||
"""
|
||||
Create a VIT instance from config.
|
||||
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
A VIT backbone instance.
|
||||
"""
|
||||
# fmt: off
|
||||
name = cfg.MODEL.VIT.NAME
|
||||
out_features = cfg.MODEL.VIT.OUT_FEATURES
|
||||
drop_path = cfg.MODEL.VIT.DROP_PATH
|
||||
img_size = cfg.MODEL.VIT.IMG_SIZE
|
||||
pos_type = cfg.MODEL.VIT.POS_TYPE
|
||||
|
||||
model_kwargs = eval(str(cfg.MODEL.VIT.MODEL_KWARGS).replace("`", ""))
|
||||
|
||||
return VIT_Backbone(name, out_features, drop_path, img_size, pos_type, model_kwargs)
|
||||
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def build_vit_fpn_backbone(cfg, input_shape: ShapeSpec):
|
||||
"""
|
||||
Create a VIT w/ FPN backbone.
|
||||
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_VIT_backbone(cfg)
|
||||
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm=cfg.MODEL.FPN.NORM,
|
||||
top_block=LastLevelMaxPool(),
|
||||
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||
)
|
||||
return backbone
|
||||
@@ -0,0 +1,687 @@
|
||||
""" 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 warnings
|
||||
import math
|
||||
import torch
|
||||
from functools import partial
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint as checkpoint
|
||||
from timm.models.layers import drop_path, to_2tuple, trunc_normal_
|
||||
|
||||
|
||||
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 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.):
|
||||
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)
|
||||
# commit this for the orignal BERT implement
|
||||
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., window_size=None, attn_head_dim=None):
|
||||
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
|
||||
# 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, 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
|
||||
|
||||
if window_size:
|
||||
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)
|
||||
|
||||
# trunc_normal_(self.relative_position_bias_table, std=.0)
|
||||
else:
|
||||
self.window_size = None
|
||||
self.relative_position_bias_table = None
|
||||
self.relative_position_index = None
|
||||
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj = nn.Linear(all_head_dim, dim)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
|
||||
def cogview_attn(self, attention_scores, alpha=32):
|
||||
'''
|
||||
https://arxiv.org/pdf/2105.13290.pdf
|
||||
Section 2.4 Stabilization of training: Precision Bottleneck Relaxation (PB-Relax).
|
||||
A replacement of the original nn.Softmax(dim=-1)(attention_scores)
|
||||
Seems the new attention_probs will result in a slower speed and a little bias
|
||||
Can use torch.allclose(standard_attention_probs, cogview_attention_probs, atol=1e-08) for comparison
|
||||
The smaller atol (e.g., 1e-08), the better.
|
||||
'''
|
||||
scaled_attention_scores = attention_scores / alpha
|
||||
max_value = scaled_attention_scores.amax(dim=(-1)).unsqueeze(-1)
|
||||
# max_value = scaled_attention_scores.amax(dim=(-2, -1)).unsqueeze(-1).unsqueeze(-1)
|
||||
new_attention_scores = (scaled_attention_scores - max_value) * alpha
|
||||
return nn.Softmax(dim=-1)(new_attention_scores)
|
||||
|
||||
def forward(self, x, rel_pos_bias=None, training_window_size=None):
|
||||
B, N, C = x.shape
|
||||
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 = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
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)
|
||||
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 self.relative_position_bias_table is not None:
|
||||
if training_window_size == self.window_size:
|
||||
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)
|
||||
else:
|
||||
training_window_size = tuple(training_window_size.tolist())
|
||||
new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
|
||||
# new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
|
||||
new_relative_position_bias_table = F.interpolate(
|
||||
self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
|
||||
2 * self.window_size[0] - 1,
|
||||
2 * self.window_size[1] - 1),
|
||||
size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
|
||||
align_corners=False)
|
||||
new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
|
||||
new_num_relative_distance - 3).permute(
|
||||
1, 0)
|
||||
new_relative_position_bias_table = torch.cat(
|
||||
[new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(training_window_size[0])
|
||||
coords_w = torch.arange(training_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] += training_window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += training_window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(training_window_size[0] * training_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:] = new_num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = new_num_relative_distance - 2
|
||||
relative_position_index[0, 0] = new_num_relative_distance - 1
|
||||
|
||||
relative_position_bias = \
|
||||
new_relative_position_bias_table[relative_position_index.view(-1)].view(
|
||||
training_window_size[0] * training_window_size[1] + 1,
|
||||
training_window_size[0] * training_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)
|
||||
|
||||
if rel_pos_bias is not None:
|
||||
attn = attn + rel_pos_bias
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
# attn = self.cogview_attn(attn)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
|
||||
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., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
|
||||
window_size=None, attn_head_dim=None):
|
||||
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, window_size=window_size, attn_head_dim=attn_head_dim)
|
||||
# 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)
|
||||
|
||||
if init_values is not None:
|
||||
self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
|
||||
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, rel_pos_bias=None, training_window_size=None):
|
||||
if self.gamma_1 is None:
|
||||
x = x + self.drop_path(
|
||||
self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, training_window_size=training_window_size))
|
||||
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
||||
else:
|
||||
x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias,
|
||||
training_window_size=training_window_size))
|
||||
x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
""" Image to Patch Embedding
|
||||
"""
|
||||
|
||||
def __init__(self, img_size=[224, 224], patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
|
||||
self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
self.num_patches_w = self.patch_shape[0]
|
||||
self.num_patches_h = self.patch_shape[1]
|
||||
# the so-called patch_shape is the patch shape during pre-training
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.num_patches = num_patches
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x, position_embedding=None, **kwargs):
|
||||
# FIXME look at relaxing size constraints
|
||||
# assert H == self.img_size[0] and W == self.img_size[1], \
|
||||
# f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
||||
x = self.proj(x)
|
||||
Hp, Wp = x.shape[2], x.shape[3]
|
||||
|
||||
if position_embedding is not None:
|
||||
# interpolate the position embedding to the corresponding size
|
||||
position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3,
|
||||
1, 2)
|
||||
position_embedding = F.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
|
||||
x = x + position_embedding
|
||||
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
return x, (Hp, Wp)
|
||||
|
||||
|
||||
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, 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 RelativePositionBias(nn.Module):
|
||||
|
||||
def __init__(self, window_size, num_heads):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.num_heads = num_heads
|
||||
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)
|
||||
|
||||
# trunc_normal_(self.relative_position_bias_table, std=.02)
|
||||
|
||||
def forward(self, training_window_size):
|
||||
if training_window_size == self.window_size:
|
||||
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
|
||||
else:
|
||||
training_window_size = tuple(training_window_size.tolist())
|
||||
new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
|
||||
# new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
|
||||
new_relative_position_bias_table = F.interpolate(
|
||||
self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
|
||||
2 * self.window_size[0] - 1,
|
||||
2 * self.window_size[1] - 1),
|
||||
size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
|
||||
align_corners=False)
|
||||
new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
|
||||
new_num_relative_distance - 3).permute(
|
||||
1, 0)
|
||||
new_relative_position_bias_table = torch.cat(
|
||||
[new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(training_window_size[0])
|
||||
coords_w = torch.arange(training_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] += training_window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += training_window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(training_window_size[0] * training_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:] = new_num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = new_num_relative_distance - 2
|
||||
relative_position_index[0, 0] = new_num_relative_distance - 1
|
||||
|
||||
relative_position_bias = \
|
||||
new_relative_position_bias_table[relative_position_index.view(-1)].view(
|
||||
training_window_size[0] * training_window_size[1] + 1,
|
||||
training_window_size[0] * training_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
|
||||
|
||||
return relative_position_bias
|
||||
|
||||
|
||||
class BEiT(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
img_size=[224, 224],
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
num_classes=80,
|
||||
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=None,
|
||||
init_values=None,
|
||||
use_abs_pos_emb=False,
|
||||
use_rel_pos_bias=False,
|
||||
use_shared_rel_pos_bias=False,
|
||||
use_checkpoint=True,
|
||||
pretrained=None,
|
||||
out_features=None,
|
||||
):
|
||||
|
||||
super(BEiT, self).__init__()
|
||||
|
||||
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
|
||||
self.num_classes = num_classes
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
if hybrid_backbone is not None:
|
||||
self.patch_embed = HybridEmbed(
|
||||
hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
else:
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
self.out_features = out_features
|
||||
self.out_indices = [int(name[5:]) for name in out_features]
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
# self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
if use_abs_pos_emb:
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
|
||||
else:
|
||||
self.pos_embed = None
|
||||
self.pos_drop = nn.Dropout(p=drop_rate)
|
||||
|
||||
self.use_shared_rel_pos_bias = use_shared_rel_pos_bias
|
||||
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
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
self.use_rel_pos_bias = use_rel_pos_bias
|
||||
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,
|
||||
init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)
|
||||
for i in range(depth)])
|
||||
|
||||
# trunc_normal_(self.mask_token, std=.02)
|
||||
|
||||
if patch_size == 16:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
# nn.SyncBatchNorm(embed_dim),
|
||||
nn.BatchNorm2d(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn3 = nn.Identity()
|
||||
|
||||
self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
elif patch_size == 8:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Identity()
|
||||
|
||||
self.fpn3 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn4 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=4, stride=4),
|
||||
)
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
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 init_weights(self):
|
||||
"""Initialize the weights in backbone.
|
||||
|
||||
Args:
|
||||
pretrained (str, optional): Path to pre-trained weights.
|
||||
Defaults to None.
|
||||
"""
|
||||
logger = get_root_logger()
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
if self.init_cfg is None:
|
||||
logger.warn(f'No pre-trained weights for '
|
||||
f'{self.__class__.__name__}, '
|
||||
f'training start from scratch')
|
||||
else:
|
||||
assert 'checkpoint' in self.init_cfg, f'Only support ' \
|
||||
f'specify `Pretrained` in ' \
|
||||
f'`init_cfg` in ' \
|
||||
f'{self.__class__.__name__} '
|
||||
logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
|
||||
load_checkpoint(self,
|
||||
filename=self.init_cfg['checkpoint'],
|
||||
strict=False,
|
||||
logger=logger,
|
||||
beit_spec_expand_rel_pos = self.use_rel_pos_bias,
|
||||
)
|
||||
'''
|
||||
|
||||
def get_num_layers(self):
|
||||
return len(self.blocks)
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'pos_embed', 'cls_token'}
|
||||
|
||||
def forward_features(self, x):
|
||||
B, C, H, W = x.shape
|
||||
x, (Hp, Wp) = self.patch_embed(x, self.pos_embed[:, 1:, :] if self.pos_embed is not None else None)
|
||||
# Hp, Wp are HW for patches
|
||||
batch_size, seq_len, _ = x.size()
|
||||
|
||||
cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
if self.pos_embed is not None:
|
||||
cls_tokens = cls_tokens + self.pos_embed[:, :1, :]
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
x = self.pos_drop(x)
|
||||
|
||||
features = []
|
||||
training_window_size = torch.tensor([Hp, Wp])
|
||||
|
||||
rel_pos_bias = self.rel_pos_bias(training_window_size) if self.rel_pos_bias is not None else None
|
||||
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if self.use_checkpoint:
|
||||
x = checkpoint.checkpoint(blk, x, rel_pos_bias, training_window_size)
|
||||
else:
|
||||
x = blk(x, rel_pos_bias=rel_pos_bias, training_window_size=training_window_size)
|
||||
if i in self.out_indices:
|
||||
xp = x[:, 1:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
|
||||
features.append(xp.contiguous())
|
||||
|
||||
ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
|
||||
for i in range(len(features)):
|
||||
features[i] = ops[i](features[i])
|
||||
|
||||
feat_out = {}
|
||||
|
||||
for name, value in zip(self.out_features, features):
|
||||
feat_out[name] = value
|
||||
|
||||
return feat_out
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
return x
|
||||
|
||||
|
||||
def beit_base_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=None,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def beit_large_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=None,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def dit_base_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=0.1,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def dit_large_patch16(pretrained=False, **kwargs):
|
||||
model = BEiT(
|
||||
patch_size=16,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
init_values=1e-5,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = BEiT(use_checkpoint=True, use_shared_rel_pos_bias=True)
|
||||
model = model.to("cuda:0")
|
||||
input1 = torch.rand(2, 3, 512, 762).to("cuda:0")
|
||||
input2 = torch.rand(2, 3, 800, 1200).to("cuda:0")
|
||||
input3 = torch.rand(2, 3, 720, 1000).to("cuda:0")
|
||||
output1 = model(input1)
|
||||
output2 = model(input2)
|
||||
output3 = model(input3)
|
||||
print("all done")
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
# File : __init__.py
|
||||
# Author : Zhaoyi Wan <wanzhaoyi@megvii.com>
|
||||
# Date : 21.11.2018
|
||||
# Last Modified Date: 08.01.2019
|
||||
# Last Modified By : Zhaoyi Wan <wanzhaoyi@megvii.com>
|
||||
|
||||
from .log import Logger
|
||||
from .average_meter import AverageMeter
|
||||
from .visualizer import Visualize
|
||||
from .box2seg import resize_with_coordinates, box2seg
|
||||
from .convert import convert
|
||||
@@ -0,0 +1,17 @@
|
||||
class AverageMeter(object):
|
||||
"""Computes and stores the average and current value"""
|
||||
def __init__(self):
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.val = 0
|
||||
self.avg = 0
|
||||
self.sum = 0
|
||||
self.count = 0
|
||||
|
||||
def update(self, val, n=1):
|
||||
self.val = val
|
||||
self.sum += val * n
|
||||
self.count += n
|
||||
self.avg = self.sum / self.count
|
||||
return self
|
||||
@@ -0,0 +1,69 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
|
||||
def intersection(x, p1, p2):
|
||||
x1, y1 = p1
|
||||
x2, y2 = p2
|
||||
if x2 == x1:
|
||||
return 0
|
||||
k = (x - x1) / (x2 - x1)
|
||||
return k * (y2 - y1) + y1
|
||||
|
||||
|
||||
def midpoint(p1, p2, typed=float):
|
||||
return [typed((p1[0] + p2[0]) / 2), typed((p1[1] + p2[1]) / 2)]
|
||||
|
||||
|
||||
def resize_with_coordinates(image, width, height, coordinates):
|
||||
original_height, original_width = image.shape[:2]
|
||||
resized_image = cv2.resize(image, (width, height))
|
||||
if coordinates is not None:
|
||||
assert coordinates.ndim == 2
|
||||
assert coordinates.shape[-1] == 2
|
||||
|
||||
rate_x = width / original_width
|
||||
rate_y = height / original_height
|
||||
|
||||
coordinates = coordinates * (rate_x, rate_y)
|
||||
return resized_image, coordinates
|
||||
|
||||
|
||||
def box2seg(image, boxes, label):
|
||||
height, width = image.shape[:2]
|
||||
mask = np.zeros((height, width), dtype=np.float32)
|
||||
seg = np.zeros((height, width), dtype=np.float32)
|
||||
points = []
|
||||
for box_index in range(boxes.shape[0]):
|
||||
box = boxes[box_index, :, :] # 4x2
|
||||
left_top = box[0]
|
||||
right_top = box[1]
|
||||
right_bottom = box[2]
|
||||
left_bottom = box[3]
|
||||
|
||||
left = [(left_top[0] + left_bottom[0]) / 2, (left_top[1] + left_bottom[1]) / 2]
|
||||
right = [(right_top[0] + right_bottom[0]) / 2, (right_top[1] + right_bottom[1]) / 2]
|
||||
|
||||
center = midpoint(left, right)
|
||||
points.append(midpoint(left, center))
|
||||
points.append(midpoint(right, center))
|
||||
|
||||
poly = np.array([midpoint(left_top, center),
|
||||
midpoint(right_top, center),
|
||||
midpoint(right_bottom, center),
|
||||
midpoint(left_bottom, center)
|
||||
])
|
||||
seg = cv2.fillPoly(seg, [poly.reshape(4, 1, 2).astype(np.int32)], int(label[box_index]))
|
||||
|
||||
left_y = intersection(0, points[0], points[1])
|
||||
right_y = intersection(width, points[-1], points[-2])
|
||||
points.insert(0, [0, left_y])
|
||||
points.append([width, right_y])
|
||||
points = np.array(points)
|
||||
|
||||
f = interpolate.interp1d(points[:, 0], points[:, 1], fill_value='extrapolate')
|
||||
xnew = np.arange(0, width, 1)
|
||||
ynew = f(xnew).clip(0, height-1)
|
||||
for x in range(width - 1):
|
||||
mask[int(ynew[x]), x] = 1
|
||||
return ynew.reshape(1, -1).round(), seg
|
||||
@@ -0,0 +1,191 @@
|
||||
import importlib
|
||||
from collections import OrderedDict
|
||||
|
||||
import anyconfig
|
||||
import munch
|
||||
|
||||
|
||||
class Config(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def load(self, conf):
|
||||
conf = anyconfig.load(conf)
|
||||
return munch.munchify(conf)
|
||||
|
||||
def compile(self, conf, return_packages=False):
|
||||
packages = conf.get('package', [])
|
||||
defines = {}
|
||||
|
||||
for path in conf.get('import', []):
|
||||
parent_conf = self.load(path)
|
||||
parent_packages, parent_defines = self.compile(
|
||||
parent_conf, return_packages=True)
|
||||
packages.extend(parent_packages)
|
||||
defines.update(parent_defines)
|
||||
|
||||
modules = []
|
||||
for package in packages:
|
||||
module = importlib.import_module(package)
|
||||
modules.append(module)
|
||||
|
||||
if isinstance(conf['define'], dict):
|
||||
conf['define'] = [conf['define']]
|
||||
|
||||
for define in conf['define']:
|
||||
name = define.copy().pop('name')
|
||||
|
||||
if not isinstance(name, str):
|
||||
raise RuntimeError('name must be str')
|
||||
|
||||
defines[name] = self.compile_conf(define, defines, modules)
|
||||
|
||||
if return_packages:
|
||||
return packages, defines
|
||||
else:
|
||||
return defines
|
||||
|
||||
def compile_conf(self, conf, defines, modules):
|
||||
if isinstance(conf, (int, float)):
|
||||
return conf
|
||||
elif isinstance(conf, str):
|
||||
if conf.startswith('^'):
|
||||
return defines[conf[1:]]
|
||||
if conf.startswith('$'):
|
||||
return {'class': self.find_class_in_modules(conf[1:], modules)}
|
||||
return conf
|
||||
elif isinstance(conf, dict):
|
||||
if 'class' in conf:
|
||||
conf['class'] = self.find_class_in_modules(
|
||||
conf['class'], modules)
|
||||
if 'base' in conf:
|
||||
base = conf.copy().pop('base')
|
||||
|
||||
if not isinstance(base, str):
|
||||
raise RuntimeError('base must be str')
|
||||
|
||||
conf = {
|
||||
**defines[base],
|
||||
**conf,
|
||||
}
|
||||
return {key: self.compile_conf(value, defines, modules) for key, value in conf.items()}
|
||||
elif isinstance(conf, (list, tuple)):
|
||||
return [self.compile_conf(value, defines, modules) for value in conf]
|
||||
else:
|
||||
return conf
|
||||
|
||||
def find_class_in_modules(self, cls, modules):
|
||||
if not isinstance(cls, str):
|
||||
raise RuntimeError('class name must be str')
|
||||
|
||||
if cls.find('.') != -1:
|
||||
package, cls = cls.rsplit('.', 1)
|
||||
module = importlib.import_module(package)
|
||||
if hasattr(module, cls):
|
||||
return module.__name__ + '.' + cls
|
||||
|
||||
for module in modules:
|
||||
if hasattr(module, cls):
|
||||
return module.__name__ + '.' + cls
|
||||
raise RuntimeError('class not found ' + cls)
|
||||
|
||||
|
||||
class State:
|
||||
def __init__(self, autoload=True, default=None):
|
||||
self.autoload = autoload
|
||||
self.default = default
|
||||
|
||||
|
||||
class StateMeta(type):
|
||||
def __new__(mcs, name, bases, attrs):
|
||||
current_states = []
|
||||
for key, value in attrs.items():
|
||||
if isinstance(value, State):
|
||||
current_states.append((key, value))
|
||||
|
||||
current_states.sort(key=lambda x: x[0])
|
||||
attrs['states'] = OrderedDict(current_states)
|
||||
new_class = super(StateMeta, mcs).__new__(mcs, name, bases, attrs)
|
||||
|
||||
# Walk through the MRO
|
||||
states = OrderedDict()
|
||||
for base in reversed(new_class.__mro__):
|
||||
if hasattr(base, 'states'):
|
||||
states.update(base.states)
|
||||
new_class.states = states
|
||||
|
||||
for key, value in states.items():
|
||||
setattr(new_class, key, value.default)
|
||||
|
||||
return new_class
|
||||
|
||||
|
||||
class Configurable(metaclass=StateMeta):
|
||||
def __init__(self, *args, cmd={}, **kwargs):
|
||||
self.load_all(cmd=cmd, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def construct_class_from_config(args):
|
||||
cls = Configurable.extract_class_from_args(args)
|
||||
return cls(**args)
|
||||
|
||||
@staticmethod
|
||||
def extract_class_from_args(args):
|
||||
cls = args.copy().pop('class')
|
||||
package, cls = cls.rsplit('.', 1)
|
||||
module = importlib.import_module(package)
|
||||
cls = getattr(module, cls)
|
||||
return cls
|
||||
|
||||
def load_all(self, **kwargs):
|
||||
for name, state in self.states.items():
|
||||
if state.autoload:
|
||||
self.load(name, **kwargs)
|
||||
|
||||
def load(self, state_name, **kwargs):
|
||||
# FIXME: kwargs should be filtered
|
||||
# Args passed from command line
|
||||
cmd = kwargs.pop('cmd', dict())
|
||||
if state_name in kwargs:
|
||||
setattr(self, state_name, self.create_member_from_config(
|
||||
(kwargs[state_name], cmd)))
|
||||
else:
|
||||
setattr(self, state_name, self.states[state_name].default)
|
||||
|
||||
def create_member_from_config(self, conf):
|
||||
args, cmd = conf
|
||||
if args is None or isinstance(args, (int, float, str)):
|
||||
return args
|
||||
elif isinstance(args, (list, tuple)):
|
||||
return [self.create_member_from_config((subargs, cmd)) for subargs in args]
|
||||
elif isinstance(args, dict):
|
||||
if 'class' in args:
|
||||
cls = self.extract_class_from_args(args)
|
||||
return cls(**args, cmd=cmd)
|
||||
return {key: self.create_member_from_config((subargs, cmd)) for key, subargs in args.items()}
|
||||
else:
|
||||
return args
|
||||
|
||||
def dump(self):
|
||||
state = {}
|
||||
state['class'] = self.__class__.__module__ + \
|
||||
'.' + self.__class__.__name__
|
||||
for name, value in self.states.items():
|
||||
obj = getattr(self, name)
|
||||
state[name] = self.dump_obj(obj)
|
||||
return state
|
||||
|
||||
def dump_obj(self, obj):
|
||||
if obj is None:
|
||||
return None
|
||||
elif hasattr(obj, 'dump'):
|
||||
return obj.dump()
|
||||
elif isinstance(obj, (int, float, str)):
|
||||
return obj
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
return [self.dump_obj(value) for value in obj]
|
||||
elif isinstance(obj, dict):
|
||||
return {key: self.dump_obj(value) for key, value in obj.items()}
|
||||
else:
|
||||
return str(obj)
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import base64
|
||||
import io
|
||||
import numpy as np
|
||||
|
||||
|
||||
def convert(data):
|
||||
if isinstance(data, dict):
|
||||
ndata = {}
|
||||
for key, value in data.items():
|
||||
nkey = key.decode()
|
||||
if nkey == 'img':
|
||||
img = Image.open(io.BytesIO(value))
|
||||
img = img.convert('RGB')
|
||||
img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
|
||||
nvalue = img
|
||||
else:
|
||||
nvalue = convert(value)
|
||||
ndata[nkey] = nvalue
|
||||
return ndata
|
||||
elif isinstance(data, list):
|
||||
return [convert(item) for item in data]
|
||||
elif isinstance(data, bytes):
|
||||
return data.decode()
|
||||
else:
|
||||
return data
|
||||
|
||||
|
||||
def to_np(x):
|
||||
return x.cpu().data.numpy()
|
||||
@@ -0,0 +1,323 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
import math
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
|
||||
class DetectionDetEvalEvaluator(object):
|
||||
def __init__(
|
||||
self,
|
||||
area_recall_constraint=0.8, area_precision_constraint=0.4,
|
||||
ev_param_ind_center_diff_thr=1,
|
||||
mtype_oo_o=1.0, mtype_om_o=0.8, mtype_om_m=1.0
|
||||
):
|
||||
|
||||
|
||||
self.area_recall_constraint = area_recall_constraint
|
||||
self.area_precision_constraint = area_precision_constraint
|
||||
self.ev_param_ind_center_diff_thr = ev_param_ind_center_diff_thr
|
||||
self.mtype_oo_o = mtype_oo_o
|
||||
self.mtype_om_o = mtype_om_o
|
||||
self.mtype_om_m = mtype_om_m
|
||||
|
||||
def evaluate_image(self, gt, pred):
|
||||
|
||||
def get_union(pD,pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
def get_intersection_over_union(pD,pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
def get_intersection(pD,pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
def one_to_one_match(row, col):
|
||||
cont = 0
|
||||
for j in range(len(recallMat[0])):
|
||||
if recallMat[row,j] >= self.area_recall_constraint and precisionMat[row,j] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
cont = 0
|
||||
for i in range(len(recallMat)):
|
||||
if recallMat[i,col] >= self.area_recall_constraint and precisionMat[i,col] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
|
||||
if recallMat[row,col] >= self.area_recall_constraint and precisionMat[row,col] >= self.area_precision_constraint:
|
||||
return True
|
||||
return False
|
||||
|
||||
def num_overlaps_gt(gtNum):
|
||||
cont = 0
|
||||
for detNum in range(len(detRects)):
|
||||
if detNum not in detDontCareRectsNum:
|
||||
if recallMat[gtNum,detNum] > 0 :
|
||||
cont = cont +1
|
||||
return cont
|
||||
|
||||
def num_overlaps_det(detNum):
|
||||
cont = 0
|
||||
for gtNum in range(len(recallMat)):
|
||||
if gtNum not in gtDontCareRectsNum:
|
||||
if recallMat[gtNum,detNum] > 0 :
|
||||
cont = cont +1
|
||||
return cont
|
||||
|
||||
def is_single_overlap(row, col):
|
||||
if num_overlaps_gt(row)==1 and num_overlaps_det(col)==1:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def one_to_many_match(gtNum):
|
||||
many_sum = 0
|
||||
detRects = []
|
||||
for detNum in range(len(recallMat[0])):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and detNum not in detDontCareRectsNum:
|
||||
if precisionMat[gtNum,detNum] >= self.area_precision_constraint:
|
||||
many_sum += recallMat[gtNum,detNum]
|
||||
detRects.append(detNum)
|
||||
if round(many_sum,4) >= self.area_recall_constraint:
|
||||
return True,detRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def many_to_one_match(detNum):
|
||||
many_sum = 0
|
||||
gtRects = []
|
||||
for gtNum in range(len(recallMat)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum:
|
||||
if recallMat[gtNum,detNum] >= self.area_recall_constraint:
|
||||
many_sum += precisionMat[gtNum,detNum]
|
||||
gtRects.append(gtNum)
|
||||
if round(many_sum,4) >= self.area_precision_constraint:
|
||||
return True,gtRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def center_distance(r1, r2):
|
||||
return ((np.mean(r1, axis=0) - np.mean(r2, axis=0)) ** 2).sum() ** 0.5
|
||||
|
||||
def diag(r):
|
||||
r = np.array(r)
|
||||
return ((r[:, 0].max() - r[:, 0].min()) ** 2 + (r[:, 1].max() - r[:, 1].min()) ** 2) ** 0.5
|
||||
|
||||
perSampleMetrics = {}
|
||||
|
||||
recall = 0
|
||||
precision = 0
|
||||
hmean = 0
|
||||
recallAccum = 0.
|
||||
precisionAccum = 0.
|
||||
gtRects = []
|
||||
detRects = []
|
||||
gtPolPoints = []
|
||||
detPolPoints = []
|
||||
gtDontCareRectsNum = []#Array of Ground Truth Rectangles' keys marked as don't Care
|
||||
detDontCareRectsNum = []#Array of Detected Rectangles' matched with a don't Care GT
|
||||
pairs = []
|
||||
evaluationLog = ""
|
||||
|
||||
recallMat = np.empty([1,1])
|
||||
precisionMat = np.empty([1,1])
|
||||
|
||||
for n in range(len(gt)):
|
||||
points = gt[n]['points']
|
||||
# transcription = gt[n]['text']
|
||||
dontCare = gt[n]['ignore']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
gtRects.append(points)
|
||||
gtPolPoints.append(points)
|
||||
if dontCare:
|
||||
gtDontCareRectsNum.append( len(gtRects)-1 )
|
||||
|
||||
evaluationLog += "GT rectangles: " + str(len(gtRects)) + (" (" + str(len(gtDontCareRectsNum)) + " don't care)\n" if len(gtDontCareRectsNum)>0 else "\n")
|
||||
|
||||
for n in range(len(pred)):
|
||||
points = pred[n]['points']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
detRect = points
|
||||
detRects.append(detRect)
|
||||
detPolPoints.append(points)
|
||||
if len(gtDontCareRectsNum)>0 :
|
||||
for dontCareRectNum in gtDontCareRectsNum:
|
||||
dontCareRect = gtRects[dontCareRectNum]
|
||||
intersected_area = get_intersection(dontCareRect,detRect)
|
||||
rdDimensions = Polygon(detRect).area
|
||||
if (rdDimensions==0) :
|
||||
precision = 0
|
||||
else:
|
||||
precision= intersected_area / rdDimensions
|
||||
if (precision > self.area_precision_constraint):
|
||||
detDontCareRectsNum.append( len(detRects)-1 )
|
||||
break
|
||||
|
||||
evaluationLog += "DET rectangles: " + str(len(detRects)) + (" (" + str(len(detDontCareRectsNum)) + " don't care)\n" if len(detDontCareRectsNum)>0 else "\n")
|
||||
|
||||
if len(gtRects)==0:
|
||||
recall = 1
|
||||
precision = 0 if len(detRects)>0 else 1
|
||||
|
||||
if len(detRects)>0:
|
||||
#Calculate recall and precision matrixs
|
||||
outputShape=[len(gtRects),len(detRects)]
|
||||
recallMat = np.empty(outputShape)
|
||||
precisionMat = np.empty(outputShape)
|
||||
gtRectMat = np.zeros(len(gtRects),np.int8)
|
||||
detRectMat = np.zeros(len(detRects),np.int8)
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
intersected_area = get_intersection(rG,rD)
|
||||
rgDimensions = Polygon(rG).area
|
||||
rdDimensions = Polygon(rD).area
|
||||
recallMat[gtNum,detNum] = 0 if rgDimensions==0 else intersected_area / rgDimensions
|
||||
precisionMat[gtNum,detNum] = 0 if rdDimensions==0 else intersected_area / rdDimensions
|
||||
|
||||
# Find one-to-one matches
|
||||
evaluationLog += "Find one-to-one matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum and detNum not in detDontCareRectsNum :
|
||||
match = one_to_one_match(gtNum, detNum)
|
||||
if match is True :
|
||||
#in deteval we have to make other validation before mark as one-to-one
|
||||
if is_single_overlap(gtNum, detNum) is True :
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
normDist = center_distance(rG, rD);
|
||||
normDist /= diag(rG) + diag(rD);
|
||||
normDist *= 2.0;
|
||||
if normDist < self.ev_param_ind_center_diff_thr:
|
||||
gtRectMat[gtNum] = 1
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += self.mtype_oo_o
|
||||
precisionAccum += self.mtype_oo_o
|
||||
pairs.append({'gt':gtNum,'det':detNum,'type':'OO'})
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(detNum) + "\n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(gtNum) + " with Det #" + str(detNum) + " normDist: " + str(normDist) + " \n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(gtNum) + " with Det #" + str(detNum) + " not single overlap\n"
|
||||
# Find one-to-many matches
|
||||
evaluationLog += "Find one-to-many matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
if gtNum not in gtDontCareRectsNum:
|
||||
match,matchesDet = one_to_many_match(gtNum)
|
||||
if match is True :
|
||||
evaluationLog += "num_overlaps_gt=" + str(num_overlaps_gt(gtNum))
|
||||
#in deteval we have to make other validation before mark as one-to-one
|
||||
if num_overlaps_gt(gtNum)>=2 :
|
||||
gtRectMat[gtNum] = 1
|
||||
recallAccum += (self.mtype_oo_o if len(matchesDet)==1 else self.mtype_om_o)
|
||||
precisionAccum += (self.mtype_oo_o if len(matchesDet)==1 else self.mtype_om_o*len(matchesDet))
|
||||
pairs.append({'gt':gtNum,'det':matchesDet,'type': 'OO' if len(matchesDet)==1 else 'OM'})
|
||||
for detNum in matchesDet :
|
||||
detRectMat[detNum] = 1
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(matchesDet) + "\n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(gtNum) + " with Det #" + str(matchesDet) + " not single overlap\n"
|
||||
|
||||
# Find many-to-one matches
|
||||
evaluationLog += "Find many-to-one matches\n"
|
||||
for detNum in range(len(detRects)):
|
||||
if detNum not in detDontCareRectsNum:
|
||||
match,matchesGt = many_to_one_match(detNum)
|
||||
if match is True :
|
||||
#in deteval we have to make other validation before mark as one-to-one
|
||||
if num_overlaps_det(detNum)>=2 :
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += (self.mtype_oo_o if len(matchesGt)==1 else self.mtype_om_m*len(matchesGt))
|
||||
precisionAccum += (self.mtype_oo_o if len(matchesGt)==1 else self.mtype_om_m)
|
||||
pairs.append({'gt':matchesGt,'det':detNum,'type': 'OO' if len(matchesGt)==1 else 'MO'})
|
||||
for gtNum in matchesGt :
|
||||
gtRectMat[gtNum] = 1
|
||||
evaluationLog += "Match GT #" + str(matchesGt) + " with Det #" + str(detNum) + "\n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(matchesGt) + " with Det #" + str(detNum) + " not single overlap\n"
|
||||
|
||||
numGtCare = (len(gtRects) - len(gtDontCareRectsNum))
|
||||
if numGtCare == 0:
|
||||
recall = float(1)
|
||||
precision = float(0) if len(detRects)>0 else float(1)
|
||||
else:
|
||||
recall = float(recallAccum) / numGtCare
|
||||
precision = float(0) if (len(detRects) - len(detDontCareRectsNum))==0 else float(precisionAccum) / (len(detRects) - len(detDontCareRectsNum))
|
||||
hmean = 0 if (precision + recall)==0 else 2.0 * precision * recall / (precision + recall)
|
||||
|
||||
numGtCare = len(gtRects) - len(gtDontCareRectsNum)
|
||||
numDetCare = len(detRects) - len(detDontCareRectsNum)
|
||||
|
||||
perSampleMetrics = {
|
||||
'precision':precision,
|
||||
'recall':recall,
|
||||
'hmean':hmean,
|
||||
'pairs':pairs,
|
||||
'recallMat':[] if len(detRects)>100 else recallMat.tolist(),
|
||||
'precisionMat':[] if len(detRects)>100 else precisionMat.tolist(),
|
||||
'gtPolPoints':gtPolPoints,
|
||||
'detPolPoints':detPolPoints,
|
||||
'gtCare': numGtCare,
|
||||
'detCare': numDetCare,
|
||||
'gtDontCare':gtDontCareRectsNum,
|
||||
'detDontCare':detDontCareRectsNum,
|
||||
'recallAccum':recallAccum,
|
||||
'precisionAccum':precisionAccum,
|
||||
'evaluationLog': evaluationLog
|
||||
}
|
||||
|
||||
return perSampleMetrics
|
||||
|
||||
def combine_results(self, results):
|
||||
numGt = 0
|
||||
numDet = 0
|
||||
methodRecallSum = 0
|
||||
methodPrecisionSum = 0
|
||||
|
||||
for result in results:
|
||||
numGt += result['gtCare']
|
||||
numDet += result['detCare']
|
||||
methodRecallSum += result['recallAccum']
|
||||
methodPrecisionSum += result['precisionAccum']
|
||||
|
||||
methodRecall = 0 if numGt==0 else methodRecallSum/numGt
|
||||
methodPrecision = 0 if numDet==0 else methodPrecisionSum/numDet
|
||||
methodHmean = 0 if methodRecall + methodPrecision==0 else 2* methodRecall * methodPrecision / (methodRecall + methodPrecision)
|
||||
|
||||
methodMetrics = {'precision':methodPrecision, 'recall':methodRecall,'hmean': methodHmean }
|
||||
|
||||
return methodMetrics
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
evaluator = DetectionDetEvalEvaluator()
|
||||
gts = [[{
|
||||
'points': [(0, 0), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}, {
|
||||
'points': [(2, 2), (3, 2), (3, 3), (2, 3)],
|
||||
'text': 5678,
|
||||
'ignore': True,
|
||||
}]]
|
||||
preds = [[{
|
||||
'points': [(0.1, 0.1), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 123,
|
||||
'ignore': False,
|
||||
}]]
|
||||
results = []
|
||||
for gt, pred in zip(gts, preds):
|
||||
results.append(evaluator.evaluate_image(gt, pred))
|
||||
metrics = evaluator.combine_results(results)
|
||||
print(metrics)
|
||||
@@ -0,0 +1,290 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
import math
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
|
||||
class DetectionICDAR2013Evaluator(object):
|
||||
def __init__(
|
||||
self,
|
||||
area_recall_constraint=0.8, area_precision_constraint=0.4,
|
||||
ev_param_ind_center_diff_thr=1,
|
||||
mtype_oo_o=1.0, mtype_om_o=0.8, mtype_om_m=1.0
|
||||
):
|
||||
|
||||
|
||||
self.area_recall_constraint = area_recall_constraint
|
||||
self.area_precision_constraint = area_precision_constraint
|
||||
self.ev_param_ind_center_diff_thr = ev_param_ind_center_diff_thr
|
||||
self.mtype_oo_o = mtype_oo_o
|
||||
self.mtype_om_o = mtype_om_o
|
||||
self.mtype_om_m = mtype_om_m
|
||||
|
||||
def evaluate_image(self, gt, pred):
|
||||
|
||||
def get_union(pD,pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
def get_intersection_over_union(pD,pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
def get_intersection(pD,pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
def one_to_one_match(row, col):
|
||||
cont = 0
|
||||
for j in range(len(recallMat[0])):
|
||||
if recallMat[row,j] >= self.area_recall_constraint and precisionMat[row,j] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
cont = 0
|
||||
for i in range(len(recallMat)):
|
||||
if recallMat[i,col] >= self.area_recall_constraint and precisionMat[i,col] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
|
||||
if recallMat[row,col] >= self.area_recall_constraint and precisionMat[row,col] >= self.area_precision_constraint:
|
||||
return True
|
||||
return False
|
||||
|
||||
def one_to_many_match(gtNum):
|
||||
many_sum = 0
|
||||
detRects = []
|
||||
for detNum in range(len(recallMat[0])):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and detNum not in detDontCareRectsNum:
|
||||
if precisionMat[gtNum,detNum] >= self.area_precision_constraint:
|
||||
many_sum += recallMat[gtNum,detNum]
|
||||
detRects.append(detNum)
|
||||
if round(many_sum,4) >= self.area_recall_constraint:
|
||||
return True,detRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def many_to_one_match(detNum):
|
||||
many_sum = 0
|
||||
gtRects = []
|
||||
for gtNum in range(len(recallMat)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum:
|
||||
if recallMat[gtNum,detNum] >= self.area_recall_constraint:
|
||||
many_sum += precisionMat[gtNum,detNum]
|
||||
gtRects.append(gtNum)
|
||||
if round(many_sum,4) >= self.area_precision_constraint:
|
||||
return True,gtRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def center_distance(r1, r2):
|
||||
return ((np.mean(r1, axis=0) - np.mean(r2, axis=0)) ** 2).sum() ** 0.5
|
||||
|
||||
def diag(r):
|
||||
r = np.array(r)
|
||||
return ((r[:, 0].max() - r[:, 0].min()) ** 2 + (r[:, 1].max() - r[:, 1].min()) ** 2) ** 0.5
|
||||
|
||||
perSampleMetrics = {}
|
||||
|
||||
recall = 0
|
||||
precision = 0
|
||||
hmean = 0
|
||||
recallAccum = 0.
|
||||
precisionAccum = 0.
|
||||
gtRects = []
|
||||
detRects = []
|
||||
gtPolPoints = []
|
||||
detPolPoints = []
|
||||
gtDontCareRectsNum = []#Array of Ground Truth Rectangles' keys marked as don't Care
|
||||
detDontCareRectsNum = []#Array of Detected Rectangles' matched with a don't Care GT
|
||||
pairs = []
|
||||
evaluationLog = ""
|
||||
|
||||
recallMat = np.empty([1,1])
|
||||
precisionMat = np.empty([1,1])
|
||||
|
||||
for n in range(len(gt)):
|
||||
points = gt[n]['points']
|
||||
# transcription = gt[n]['text']
|
||||
dontCare = gt[n]['ignore']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
gtRects.append(points)
|
||||
gtPolPoints.append(points)
|
||||
if dontCare:
|
||||
gtDontCareRectsNum.append( len(gtRects)-1 )
|
||||
|
||||
evaluationLog += "GT rectangles: " + str(len(gtRects)) + (" (" + str(len(gtDontCareRectsNum)) + " don't care)\n" if len(gtDontCareRectsNum)>0 else "\n")
|
||||
|
||||
for n in range(len(pred)):
|
||||
points = pred[n]['points']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
detRect = points
|
||||
detRects.append(detRect)
|
||||
detPolPoints.append(points)
|
||||
if len(gtDontCareRectsNum)>0 :
|
||||
for dontCareRectNum in gtDontCareRectsNum:
|
||||
dontCareRect = gtRects[dontCareRectNum]
|
||||
intersected_area = get_intersection(dontCareRect,detRect)
|
||||
rdDimensions = Polygon(detRect).area
|
||||
if (rdDimensions==0) :
|
||||
precision = 0
|
||||
else:
|
||||
precision= intersected_area / rdDimensions
|
||||
if (precision > self.area_precision_constraint):
|
||||
detDontCareRectsNum.append( len(detRects)-1 )
|
||||
break
|
||||
|
||||
evaluationLog += "DET rectangles: " + str(len(detRects)) + (" (" + str(len(detDontCareRectsNum)) + " don't care)\n" if len(detDontCareRectsNum)>0 else "\n")
|
||||
|
||||
if len(gtRects)==0:
|
||||
recall = 1
|
||||
precision = 0 if len(detRects)>0 else 1
|
||||
|
||||
if len(detRects)>0:
|
||||
#Calculate recall and precision matrixs
|
||||
outputShape=[len(gtRects),len(detRects)]
|
||||
recallMat = np.empty(outputShape)
|
||||
precisionMat = np.empty(outputShape)
|
||||
gtRectMat = np.zeros(len(gtRects),np.int8)
|
||||
detRectMat = np.zeros(len(detRects),np.int8)
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
intersected_area = get_intersection(rG,rD)
|
||||
rgDimensions = Polygon(rG).area
|
||||
rdDimensions = Polygon(rD).area
|
||||
recallMat[gtNum,detNum] = 0 if rgDimensions==0 else intersected_area / rgDimensions
|
||||
precisionMat[gtNum,detNum] = 0 if rdDimensions==0 else intersected_area / rdDimensions
|
||||
|
||||
# Find one-to-one matches
|
||||
evaluationLog += "Find one-to-one matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum and detNum not in detDontCareRectsNum :
|
||||
match = one_to_one_match(gtNum, detNum)
|
||||
if match is True :
|
||||
#in deteval we have to make other validation before mark as one-to-one
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
normDist = center_distance(rG, rD);
|
||||
normDist /= diag(rG) + diag(rD);
|
||||
normDist *= 2.0;
|
||||
if normDist < self.ev_param_ind_center_diff_thr:
|
||||
gtRectMat[gtNum] = 1
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += self.mtype_oo_o
|
||||
precisionAccum += self.mtype_oo_o
|
||||
pairs.append({'gt':gtNum,'det':detNum,'type':'OO'})
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(detNum) + "\n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(gtNum) + " with Det #" + str(detNum) + " normDist: " + str(normDist) + " \n"
|
||||
# Find one-to-many matches
|
||||
evaluationLog += "Find one-to-many matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
if gtNum not in gtDontCareRectsNum:
|
||||
match,matchesDet = one_to_many_match(gtNum)
|
||||
if match is True :
|
||||
evaluationLog += "num_overlaps_gt=" + str(num_overlaps_gt(gtNum))
|
||||
gtRectMat[gtNum] = 1
|
||||
recallAccum += (self.mtype_oo_o if len(matchesDet)==1 else self.mtype_om_o)
|
||||
precisionAccum += (self.mtype_oo_o if len(matchesDet)==1 else self.mtype_om_o*len(matchesDet))
|
||||
pairs.append({'gt':gtNum,'det':matchesDet,'type': 'OO' if len(matchesDet)==1 else 'OM'})
|
||||
for detNum in matchesDet :
|
||||
detRectMat[detNum] = 1
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(matchesDet) + "\n"
|
||||
|
||||
# Find many-to-one matches
|
||||
evaluationLog += "Find many-to-one matches\n"
|
||||
for detNum in range(len(detRects)):
|
||||
if detNum not in detDontCareRectsNum:
|
||||
match,matchesGt = many_to_one_match(detNum)
|
||||
if match is True :
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += (self.mtype_oo_o if len(matchesGt)==1 else self.mtype_om_m*len(matchesGt))
|
||||
precisionAccum += (self.mtype_oo_o if len(matchesGt)==1 else self.mtype_om_m)
|
||||
pairs.append({'gt':matchesGt,'det':detNum,'type': 'OO' if len(matchesGt)==1 else 'MO'})
|
||||
for gtNum in matchesGt :
|
||||
gtRectMat[gtNum] = 1
|
||||
evaluationLog += "Match GT #" + str(matchesGt) + " with Det #" + str(detNum) + "\n"
|
||||
|
||||
numGtCare = (len(gtRects) - len(gtDontCareRectsNum))
|
||||
if numGtCare == 0:
|
||||
recall = float(1)
|
||||
precision = float(0) if len(detRects)>0 else float(1)
|
||||
else:
|
||||
recall = float(recallAccum) / numGtCare
|
||||
precision = float(0) if (len(detRects) - len(detDontCareRectsNum))==0 else float(precisionAccum) / (len(detRects) - len(detDontCareRectsNum))
|
||||
hmean = 0 if (precision + recall)==0 else 2.0 * precision * recall / (precision + recall)
|
||||
|
||||
numGtCare = len(gtRects) - len(gtDontCareRectsNum)
|
||||
numDetCare = len(detRects) - len(detDontCareRectsNum)
|
||||
|
||||
perSampleMetrics = {
|
||||
'precision':precision,
|
||||
'recall':recall,
|
||||
'hmean':hmean,
|
||||
'pairs':pairs,
|
||||
'recallMat':[] if len(detRects)>100 else recallMat.tolist(),
|
||||
'precisionMat':[] if len(detRects)>100 else precisionMat.tolist(),
|
||||
'gtPolPoints':gtPolPoints,
|
||||
'detPolPoints':detPolPoints,
|
||||
'gtCare': numGtCare,
|
||||
'detCare': numDetCare,
|
||||
'gtDontCare':gtDontCareRectsNum,
|
||||
'detDontCare':detDontCareRectsNum,
|
||||
'recallAccum':recallAccum,
|
||||
'precisionAccum':precisionAccum,
|
||||
'evaluationLog': evaluationLog
|
||||
}
|
||||
|
||||
return perSampleMetrics
|
||||
|
||||
def combine_results(self, results):
|
||||
numGt = 0
|
||||
numDet = 0
|
||||
methodRecallSum = 0
|
||||
methodPrecisionSum = 0
|
||||
|
||||
for result in results:
|
||||
numGt += result['gtCare']
|
||||
numDet += result['detCare']
|
||||
methodRecallSum += result['recallAccum']
|
||||
methodPrecisionSum += result['precisionAccum']
|
||||
|
||||
methodRecall = 0 if numGt==0 else methodRecallSum/numGt
|
||||
methodPrecision = 0 if numDet==0 else methodPrecisionSum/numDet
|
||||
methodHmean = 0 if methodRecall + methodPrecision==0 else 2* methodRecall * methodPrecision / (methodRecall + methodPrecision)
|
||||
|
||||
methodMetrics = {'precision':methodPrecision, 'recall':methodRecall,'hmean': methodHmean }
|
||||
|
||||
return methodMetrics
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
evaluator = DetectionICDAR2013Evaluator()
|
||||
gts = [[{
|
||||
'points': [(0, 0), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}, {
|
||||
'points': [(2, 2), (3, 2), (3, 3), (2, 3)],
|
||||
'text': 5678,
|
||||
'ignore': True,
|
||||
}]]
|
||||
preds = [[{
|
||||
'points': [(0.1, 0.1), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 123,
|
||||
'ignore': False,
|
||||
}]]
|
||||
results = []
|
||||
for gt, pred in zip(gts, preds):
|
||||
results.append(evaluator.evaluate_image(gt, pred))
|
||||
metrics = evaluator.combine_results(results)
|
||||
print(metrics)
|
||||
@@ -0,0 +1,222 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
|
||||
class DetectionIoUEvaluator(object):
|
||||
def __init__(self, iou_constraint=0.5, area_precision_constraint=0.5):
|
||||
self.iou_constraint = iou_constraint
|
||||
self.area_precision_constraint = area_precision_constraint
|
||||
|
||||
def evaluate_image(self, gt, pred):
|
||||
|
||||
def get_union(pD, pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
def get_intersection_over_union(pD, pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
def get_intersection(pD, pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
def compute_ap(confList, matchList, numGtCare):
|
||||
correct = 0
|
||||
AP = 0
|
||||
if len(confList) > 0:
|
||||
confList = np.array(confList)
|
||||
matchList = np.array(matchList)
|
||||
sorted_ind = np.argsort(-confList)
|
||||
confList = confList[sorted_ind]
|
||||
matchList = matchList[sorted_ind]
|
||||
for n in range(len(confList)):
|
||||
match = matchList[n]
|
||||
if match:
|
||||
correct += 1
|
||||
AP += float(correct)/(n + 1)
|
||||
|
||||
if numGtCare > 0:
|
||||
AP /= numGtCare
|
||||
|
||||
return AP
|
||||
|
||||
perSampleMetrics = {}
|
||||
|
||||
matchedSum = 0
|
||||
|
||||
Rectangle = namedtuple('Rectangle', 'xmin ymin xmax ymax')
|
||||
|
||||
numGlobalCareGt = 0
|
||||
numGlobalCareDet = 0
|
||||
|
||||
arrGlobalConfidences = []
|
||||
arrGlobalMatches = []
|
||||
|
||||
recall = 0
|
||||
precision = 0
|
||||
hmean = 0
|
||||
|
||||
detMatched = 0
|
||||
|
||||
iouMat = np.empty([1, 1])
|
||||
|
||||
gtPols = []
|
||||
detPols = []
|
||||
|
||||
gtPolPoints = []
|
||||
detPolPoints = []
|
||||
|
||||
# Array of Ground Truth Polygons' keys marked as don't Care
|
||||
gtDontCarePolsNum = []
|
||||
# Array of Detected Polygons' matched with a don't Care GT
|
||||
detDontCarePolsNum = []
|
||||
|
||||
pairs = []
|
||||
detMatchedNums = []
|
||||
|
||||
arrSampleConfidences = []
|
||||
arrSampleMatch = []
|
||||
|
||||
evaluationLog = ""
|
||||
|
||||
for n in range(len(gt)):
|
||||
points = gt[n]['points']
|
||||
# transcription = gt[n]['text']
|
||||
dontCare = gt[n]['ignore']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
gtPol = points
|
||||
gtPols.append(gtPol)
|
||||
gtPolPoints.append(points)
|
||||
if dontCare:
|
||||
gtDontCarePolsNum.append(len(gtPols)-1)
|
||||
|
||||
evaluationLog += "GT polygons: " + str(len(gtPols)) + (" (" + str(len(
|
||||
gtDontCarePolsNum)) + " don't care)\n" if len(gtDontCarePolsNum) > 0 else "\n")
|
||||
|
||||
for n in range(len(pred)):
|
||||
points = pred[n]['points']
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
detPol = points
|
||||
detPols.append(detPol)
|
||||
detPolPoints.append(points)
|
||||
if len(gtDontCarePolsNum) > 0:
|
||||
for dontCarePol in gtDontCarePolsNum:
|
||||
dontCarePol = gtPols[dontCarePol]
|
||||
intersected_area = get_intersection(dontCarePol, detPol)
|
||||
pdDimensions = Polygon(detPol).area
|
||||
precision = 0 if pdDimensions == 0 else intersected_area / pdDimensions
|
||||
if (precision > self.area_precision_constraint):
|
||||
detDontCarePolsNum.append(len(detPols)-1)
|
||||
break
|
||||
|
||||
evaluationLog += "DET polygons: " + str(len(detPols)) + (" (" + str(len(
|
||||
detDontCarePolsNum)) + " don't care)\n" if len(detDontCarePolsNum) > 0 else "\n")
|
||||
|
||||
if len(gtPols) > 0 and len(detPols) > 0:
|
||||
# Calculate IoU and precision matrixs
|
||||
outputShape = [len(gtPols), len(detPols)]
|
||||
iouMat = np.empty(outputShape)
|
||||
gtRectMat = np.zeros(len(gtPols), np.int8)
|
||||
detRectMat = np.zeros(len(detPols), np.int8)
|
||||
for gtNum in range(len(gtPols)):
|
||||
for detNum in range(len(detPols)):
|
||||
pG = gtPols[gtNum]
|
||||
pD = detPols[detNum]
|
||||
iouMat[gtNum, detNum] = get_intersection_over_union(pD, pG)
|
||||
|
||||
for gtNum in range(len(gtPols)):
|
||||
for detNum in range(len(detPols)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCarePolsNum and detNum not in detDontCarePolsNum:
|
||||
if iouMat[gtNum, detNum] > self.iou_constraint:
|
||||
gtRectMat[gtNum] = 1
|
||||
detRectMat[detNum] = 1
|
||||
detMatched += 1
|
||||
pairs.append({'gt': gtNum, 'det': detNum})
|
||||
detMatchedNums.append(detNum)
|
||||
evaluationLog += "Match GT #" + \
|
||||
str(gtNum) + " with Det #" + str(detNum) + "\n"
|
||||
|
||||
numGtCare = (len(gtPols) - len(gtDontCarePolsNum))
|
||||
numDetCare = (len(detPols) - len(detDontCarePolsNum))
|
||||
if numGtCare == 0:
|
||||
recall = float(1)
|
||||
precision = float(0) if numDetCare > 0 else float(1)
|
||||
else:
|
||||
recall = float(detMatched) / numGtCare
|
||||
precision = 0 if numDetCare == 0 else float(
|
||||
detMatched) / numDetCare
|
||||
|
||||
hmean = 0 if (precision + recall) == 0 else 2.0 * \
|
||||
precision * recall / (precision + recall)
|
||||
|
||||
matchedSum += detMatched
|
||||
numGlobalCareGt += numGtCare
|
||||
numGlobalCareDet += numDetCare
|
||||
|
||||
perSampleMetrics = {
|
||||
'precision': precision,
|
||||
'recall': recall,
|
||||
'hmean': hmean,
|
||||
'pairs': pairs,
|
||||
'iouMat': [] if len(detPols) > 100 else iouMat.tolist(),
|
||||
'gtPolPoints': gtPolPoints,
|
||||
'detPolPoints': detPolPoints,
|
||||
'gtCare': numGtCare,
|
||||
'detCare': numDetCare,
|
||||
'gtDontCare': gtDontCarePolsNum,
|
||||
'detDontCare': detDontCarePolsNum,
|
||||
'detMatched': detMatched,
|
||||
'evaluationLog': evaluationLog
|
||||
}
|
||||
|
||||
return perSampleMetrics
|
||||
|
||||
def combine_results(self, results):
|
||||
numGlobalCareGt = 0
|
||||
numGlobalCareDet = 0
|
||||
matchedSum = 0
|
||||
for result in results:
|
||||
numGlobalCareGt += result['gtCare']
|
||||
numGlobalCareDet += result['detCare']
|
||||
matchedSum += result['detMatched']
|
||||
|
||||
methodRecall = 0 if numGlobalCareGt == 0 else float(
|
||||
matchedSum)/numGlobalCareGt
|
||||
methodPrecision = 0 if numGlobalCareDet == 0 else float(
|
||||
matchedSum)/numGlobalCareDet
|
||||
methodHmean = 0 if methodRecall + methodPrecision == 0 else 2 * \
|
||||
methodRecall * methodPrecision / (methodRecall + methodPrecision)
|
||||
|
||||
methodMetrics = {'precision': methodPrecision,
|
||||
'recall': methodRecall, 'hmean': methodHmean}
|
||||
|
||||
return methodMetrics
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
evaluator = DetectionIoUEvaluator()
|
||||
gts = [[{
|
||||
'points': [(0, 0), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}, {
|
||||
'points': [(2, 2), (3, 2), (3, 3), (2, 3)],
|
||||
'text': 5678,
|
||||
'ignore': False,
|
||||
}]]
|
||||
preds = [[{
|
||||
'points': [(0.1, 0.1), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 123,
|
||||
'ignore': False,
|
||||
}]]
|
||||
results = []
|
||||
for gt, pred in zip(gts, preds):
|
||||
results.append(evaluator.evaluate_image(gt, pred))
|
||||
metrics = evaluator.combine_results(results)
|
||||
print(metrics)
|
||||
@@ -0,0 +1,285 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
import math
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
|
||||
class DetectionMTWI2018Evaluator(object):
|
||||
def __init__(
|
||||
self,
|
||||
area_recall_constraint=0.7, area_precision_constraint=0.7,
|
||||
ev_param_ind_center_diff_thr=1,
|
||||
):
|
||||
|
||||
|
||||
self.area_recall_constraint = area_recall_constraint
|
||||
self.area_precision_constraint = area_precision_constraint
|
||||
self.ev_param_ind_center_diff_thr = ev_param_ind_center_diff_thr
|
||||
|
||||
def evaluate_image(self, gt, pred):
|
||||
|
||||
def get_union(pD,pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
def get_intersection_over_union(pD,pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
def get_intersection(pD,pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
def one_to_one_match(row, col):
|
||||
cont = 0
|
||||
for j in range(len(recallMat[0])):
|
||||
if recallMat[row,j] >= self.area_recall_constraint and precisionMat[row,j] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
cont = 0
|
||||
for i in range(len(recallMat)):
|
||||
if recallMat[i,col] >= self.area_recall_constraint and precisionMat[i,col] >= self.area_precision_constraint:
|
||||
cont = cont +1
|
||||
if (cont != 1):
|
||||
return False
|
||||
|
||||
if recallMat[row,col] >= self.area_recall_constraint and precisionMat[row,col] >= self.area_precision_constraint:
|
||||
return True
|
||||
return False
|
||||
|
||||
def one_to_many_match(gtNum):
|
||||
many_sum = 0
|
||||
detRects = []
|
||||
for detNum in range(len(recallMat[0])):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and detNum not in detDontCareRectsNum:
|
||||
if precisionMat[gtNum,detNum] >= self.area_precision_constraint:
|
||||
many_sum += recallMat[gtNum,detNum]
|
||||
detRects.append(detNum)
|
||||
if round(many_sum,4) >= self.area_recall_constraint:
|
||||
return True,detRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def many_to_one_match(detNum):
|
||||
many_sum = 0
|
||||
gtRects = []
|
||||
for gtNum in range(len(recallMat)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum:
|
||||
if recallMat[gtNum,detNum] >= self.area_recall_constraint:
|
||||
many_sum += precisionMat[gtNum,detNum]
|
||||
gtRects.append(gtNum)
|
||||
if round(many_sum,4) >= self.area_precision_constraint:
|
||||
return True,gtRects
|
||||
else:
|
||||
return False,[]
|
||||
|
||||
def center_distance(r1, r2):
|
||||
return ((np.mean(r1, axis=0) - np.mean(r2, axis=0)) ** 2).sum() ** 0.5
|
||||
|
||||
def diag(r):
|
||||
r = np.array(r)
|
||||
return ((r[:, 0].max() - r[:, 0].min()) ** 2 + (r[:, 1].max() - r[:, 1].min()) ** 2) ** 0.5
|
||||
|
||||
perSampleMetrics = {}
|
||||
|
||||
recall = 0
|
||||
precision = 0
|
||||
hmean = 0
|
||||
recallAccum = 0.
|
||||
precisionAccum = 0.
|
||||
gtRects = []
|
||||
detRects = []
|
||||
gtPolPoints = []
|
||||
detPolPoints = []
|
||||
gtDontCareRectsNum = []#Array of Ground Truth Rectangles' keys marked as don't Care
|
||||
detDontCareRectsNum = []#Array of Detected Rectangles' matched with a don't Care GT
|
||||
pairs = []
|
||||
evaluationLog = ""
|
||||
|
||||
recallMat = np.empty([1,1])
|
||||
precisionMat = np.empty([1,1])
|
||||
|
||||
for n in range(len(gt)):
|
||||
points = gt[n]['points']
|
||||
# transcription = gt[n]['text']
|
||||
dontCare = gt[n]['ignore']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
gtRects.append(points)
|
||||
gtPolPoints.append(points)
|
||||
if dontCare:
|
||||
gtDontCareRectsNum.append( len(gtRects)-1 )
|
||||
|
||||
evaluationLog += "GT rectangles: " + str(len(gtRects)) + (" (" + str(len(gtDontCareRectsNum)) + " don't care)\n" if len(gtDontCareRectsNum)>0 else "\n")
|
||||
|
||||
for n in range(len(pred)):
|
||||
points = pred[n]['points']
|
||||
|
||||
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
||||
continue
|
||||
|
||||
detRect = points
|
||||
detRects.append(detRect)
|
||||
detPolPoints.append(points)
|
||||
if len(gtDontCareRectsNum)>0 :
|
||||
for dontCareRectNum in gtDontCareRectsNum:
|
||||
dontCareRect = gtRects[dontCareRectNum]
|
||||
intersected_area = get_intersection(dontCareRect,detRect)
|
||||
rdDimensions = Polygon(detRect).area
|
||||
if (rdDimensions==0) :
|
||||
precision = 0
|
||||
else:
|
||||
precision= intersected_area / rdDimensions
|
||||
if (precision > 0.5):
|
||||
detDontCareRectsNum.append( len(detRects)-1 )
|
||||
break
|
||||
|
||||
evaluationLog += "DET rectangles: " + str(len(detRects)) + (" (" + str(len(detDontCareRectsNum)) + " don't care)\n" if len(detDontCareRectsNum)>0 else "\n")
|
||||
|
||||
if len(gtRects)==0:
|
||||
recall = 1
|
||||
precision = 0 if len(detRects)>0 else 1
|
||||
|
||||
if len(detRects)>0:
|
||||
#Calculate recall and precision matrixs
|
||||
outputShape=[len(gtRects),len(detRects)]
|
||||
recallMat = np.empty(outputShape)
|
||||
precisionMat = np.empty(outputShape)
|
||||
gtRectMat = np.zeros(len(gtRects),np.int8)
|
||||
detRectMat = np.zeros(len(detRects),np.int8)
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
intersected_area = get_intersection(rG,rD)
|
||||
rgDimensions = Polygon(rG).area
|
||||
rdDimensions = Polygon(rD).area
|
||||
recallMat[gtNum,detNum] = 0 if rgDimensions==0 else intersected_area / rgDimensions
|
||||
precisionMat[gtNum,detNum] = 0 if rdDimensions==0 else intersected_area / rdDimensions
|
||||
|
||||
# Find one-to-one matches
|
||||
evaluationLog += "Find one-to-one matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
for detNum in range(len(detRects)):
|
||||
if gtRectMat[gtNum] == 0 and detRectMat[detNum] == 0 and gtNum not in gtDontCareRectsNum and detNum not in detDontCareRectsNum :
|
||||
match = one_to_one_match(gtNum, detNum)
|
||||
if match is True :
|
||||
#in deteval we have to make other validation before mark as one-to-one
|
||||
rG = gtRects[gtNum]
|
||||
rD = detRects[detNum]
|
||||
normDist = center_distance(rG, rD);
|
||||
normDist /= diag(rG) + diag(rD);
|
||||
normDist *= 2.0;
|
||||
if normDist < self.ev_param_ind_center_diff_thr:
|
||||
gtRectMat[gtNum] = 1
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += 1.0
|
||||
precisionAccum += 1.0
|
||||
pairs.append({'gt':gtNum,'det':detNum,'type':'OO'})
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(detNum) + "\n"
|
||||
else:
|
||||
evaluationLog += "Match Discarded GT #" + str(gtNum) + " with Det #" + str(detNum) + " normDist: " + str(normDist) + " \n"
|
||||
# Find one-to-many matches
|
||||
evaluationLog += "Find one-to-many matches\n"
|
||||
for gtNum in range(len(gtRects)):
|
||||
if gtNum not in gtDontCareRectsNum:
|
||||
match,matchesDet = one_to_many_match(gtNum)
|
||||
if match is True :
|
||||
gtRectMat[gtNum] = 1
|
||||
recallAccum += 1.0
|
||||
precisionAccum += len(matchesDet) / (1 + math.log(len(matchesDet)))
|
||||
pairs.append({'gt':gtNum,'det':matchesDet,'type': 'OO' if len(matchesDet)==1 else 'OM'})
|
||||
for detNum in matchesDet :
|
||||
detRectMat[detNum] = 1
|
||||
evaluationLog += "Match GT #" + str(gtNum) + " with Det #" + str(matchesDet) + "\n"
|
||||
|
||||
# Find many-to-one matches
|
||||
evaluationLog += "Find many-to-one matches\n"
|
||||
for detNum in range(len(detRects)):
|
||||
if detNum not in detDontCareRectsNum:
|
||||
match,matchesGt = many_to_one_match(detNum)
|
||||
if match is True :
|
||||
detRectMat[detNum] = 1
|
||||
recallAccum += len(matchesGt) / (1 + math.log(len(matchesGt)))
|
||||
precisionAccum += 1.0
|
||||
pairs.append({'gt':matchesGt,'det':detNum,'type': 'OO' if len(matchesGt)==1 else 'MO'})
|
||||
for gtNum in matchesGt :
|
||||
gtRectMat[gtNum] = 1
|
||||
evaluationLog += "Match GT #" + str(matchesGt) + " with Det #" + str(detNum) + "\n"
|
||||
|
||||
numGtCare = (len(gtRects) - len(gtDontCareRectsNum))
|
||||
if numGtCare == 0:
|
||||
recall = float(1)
|
||||
precision = float(0) if len(detRects)>0 else float(1)
|
||||
else:
|
||||
recall = float(recallAccum) / numGtCare
|
||||
precision = float(0) if (len(detRects) - len(detDontCareRectsNum))==0 else float(precisionAccum) / (len(detRects) - len(detDontCareRectsNum))
|
||||
hmean = 0 if (precision + recall)==0 else 2.0 * precision * recall / (precision + recall)
|
||||
|
||||
numGtCare = len(gtRects) - len(gtDontCareRectsNum)
|
||||
numDetCare = len(detRects) - len(detDontCareRectsNum)
|
||||
|
||||
perSampleMetrics = {
|
||||
'precision':precision,
|
||||
'recall':recall,
|
||||
'hmean':hmean,
|
||||
'pairs':pairs,
|
||||
'recallMat':[] if len(detRects)>100 else recallMat.tolist(),
|
||||
'precisionMat':[] if len(detRects)>100 else precisionMat.tolist(),
|
||||
'gtPolPoints':gtPolPoints,
|
||||
'detPolPoints':detPolPoints,
|
||||
'gtCare': numGtCare,
|
||||
'detCare': numDetCare,
|
||||
'gtDontCare':gtDontCareRectsNum,
|
||||
'detDontCare':detDontCareRectsNum,
|
||||
'recallAccum':recallAccum,
|
||||
'precisionAccum':precisionAccum,
|
||||
'evaluationLog': evaluationLog
|
||||
}
|
||||
|
||||
return perSampleMetrics
|
||||
|
||||
def combine_results(self, results):
|
||||
numGt = 0
|
||||
numDet = 0
|
||||
methodRecallSum = 0
|
||||
methodPrecisionSum = 0
|
||||
|
||||
for result in results:
|
||||
numGt += result['gtCare']
|
||||
numDet += result['detCare']
|
||||
methodRecallSum += result['recallAccum']
|
||||
methodPrecisionSum += result['precisionAccum']
|
||||
|
||||
methodRecall = 0 if numGt==0 else methodRecallSum/numGt
|
||||
methodPrecision = 0 if numDet==0 else methodPrecisionSum/numDet
|
||||
methodHmean = 0 if methodRecall + methodPrecision==0 else 2* methodRecall * methodPrecision / (methodRecall + methodPrecision)
|
||||
|
||||
methodMetrics = {'precision':methodPrecision, 'recall':methodRecall,'hmean': methodHmean }
|
||||
|
||||
return methodMetrics
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
evaluator = DetectionICDAR2013Evaluator()
|
||||
gts = [[{
|
||||
'points': [(0, 0), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}, {
|
||||
'points': [(2, 2), (3, 2), (3, 3), (2, 3)],
|
||||
'text': 5678,
|
||||
'ignore': True,
|
||||
}]]
|
||||
preds = [[{
|
||||
'points': [(0.1, 0.1), (1, 0), (1, 1), (0, 1)],
|
||||
'text': 123,
|
||||
'ignore': False,
|
||||
}]]
|
||||
results = []
|
||||
for gt, pred in zip(gts, preds):
|
||||
results.append(evaluator.evaluate_image(gt, pred))
|
||||
metrics = evaluator.combine_results(results)
|
||||
print(metrics)
|
||||
@@ -0,0 +1,196 @@
|
||||
import os
|
||||
import logging
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
# from tensorboardX import SummaryWriter
|
||||
import yaml
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .config import Configurable, State
|
||||
|
||||
|
||||
class Logger(Configurable):
|
||||
SUMMARY_DIR_NAME = 'summaries'
|
||||
VISUALIZE_NAME = 'visualize'
|
||||
LOG_FILE_NAME = 'output.log'
|
||||
ARGS_FILE_NAME = 'args.log'
|
||||
METRICS_FILE_NAME = 'metrics.log'
|
||||
|
||||
database_dir = State(default='./outputs/')
|
||||
log_dir = State(default='workspace')
|
||||
verbose = State(default=False)
|
||||
level = State(default='info')
|
||||
log_interval = State(default=100)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.load_all(**kwargs)
|
||||
|
||||
self._make_storage()
|
||||
|
||||
cmd = kwargs['cmd']
|
||||
self.name = cmd['name']
|
||||
self.log_dir = os.path.join(self.log_dir, self.name)
|
||||
try:
|
||||
self.verbose = cmd['verbose']
|
||||
except:
|
||||
print('verbose:', self.verbose)
|
||||
if self.verbose:
|
||||
print('Initializing log dir for', self.log_dir)
|
||||
|
||||
if not os.path.exists(self.log_dir):
|
||||
os.makedirs(self.log_dir)
|
||||
|
||||
self.message_logger = self._init_message_logger()
|
||||
|
||||
summary_path = os.path.join(self.log_dir, self.SUMMARY_DIR_NAME)
|
||||
self.tf_board_logger = SummaryWriter(summary_path)
|
||||
|
||||
self.metrics_writer = open(os.path.join(
|
||||
self.log_dir, self.METRICS_FILE_NAME), 'at')
|
||||
|
||||
self.timestamp = time.time()
|
||||
self.logged = -1
|
||||
self.speed = None
|
||||
self.eta_time = None
|
||||
|
||||
def _make_storage(self):
|
||||
application = os.path.basename(os.getcwd())
|
||||
storage_dir = os.path.join(
|
||||
self.database_dir, self.log_dir, application)
|
||||
if not os.path.exists(storage_dir):
|
||||
os.makedirs(storage_dir)
|
||||
if not os.path.exists(self.log_dir):
|
||||
os.symlink(storage_dir, self.log_dir)
|
||||
|
||||
def save_dir(self, dir_name):
|
||||
return os.path.join(self.log_dir, dir_name)
|
||||
|
||||
def _init_message_logger(self):
|
||||
message_logger = logging.getLogger('messages')
|
||||
message_logger.setLevel(
|
||||
logging.DEBUG if self.verbose else logging.INFO)
|
||||
formatter = logging.Formatter(
|
||||
'[%(levelname)s] [%(asctime)s] %(message)s')
|
||||
std_handler = logging.StreamHandler()
|
||||
std_handler.setLevel(message_logger.level)
|
||||
std_handler.setFormatter(formatter)
|
||||
|
||||
file_handler = logging.FileHandler(
|
||||
os.path.join(self.log_dir, self.LOG_FILE_NAME))
|
||||
file_handler.setLevel(message_logger.level)
|
||||
file_handler.setFormatter(formatter)
|
||||
|
||||
message_logger.addHandler(std_handler)
|
||||
message_logger.addHandler(file_handler)
|
||||
return message_logger
|
||||
|
||||
def report_time(self, name: str):
|
||||
if self.verbose:
|
||||
self.info(name + " time :" + str(time.time() - self.timestamp))
|
||||
self.timestamp = time.time()
|
||||
|
||||
def report_eta(self, steps, total, epoch):
|
||||
self.logged = self.logged % total + 1
|
||||
steps = steps % total
|
||||
if self.eta_time is None:
|
||||
self.eta_time = time.time()
|
||||
speed = -1
|
||||
else:
|
||||
eta_time = time.time()
|
||||
speed = eta_time - self.eta_time
|
||||
if self.speed is not None:
|
||||
speed = ((self.logged - 1) * self.speed + speed) / self.logged
|
||||
self.speed = speed
|
||||
self.eta_time = eta_time
|
||||
|
||||
seconds = (total - steps) * speed
|
||||
hours = seconds // 3600
|
||||
minutes = (seconds - (hours * 3600)) // 60
|
||||
seconds = seconds % 60
|
||||
|
||||
print('%d/%d batches processed in epoch %d, ETA: %2d:%2d:%2d' %
|
||||
(steps, total, epoch,
|
||||
hours, minutes, seconds), end='\r')
|
||||
|
||||
def args(self, parameters=None):
|
||||
if parameters is None:
|
||||
with open(os.path.join(self.log_dir, self.ARGS_FILE_NAME), 'rt') as reader:
|
||||
return yaml.load(reader.read())
|
||||
with open(os.path.join(self.log_dir, self.ARGS_FILE_NAME), 'wt') as writer:
|
||||
yaml.dump(parameters.dump(), writer)
|
||||
|
||||
def metrics(self, epoch, steps, metrics_dict):
|
||||
results = {}
|
||||
for name, a in metrics_dict.items():
|
||||
results[name] = {'count': a.count, 'value': float(a.avg)}
|
||||
self.add_scalar('metrics/' + name, a.avg, steps)
|
||||
result_dict = {
|
||||
str(datetime.now()): {
|
||||
'epoch': epoch,
|
||||
'steps': steps,
|
||||
**results
|
||||
}
|
||||
}
|
||||
string_result = yaml.dump(result_dict)
|
||||
self.info(string_result)
|
||||
self.metrics_writer.write(string_result)
|
||||
self.metrics_writer.flush()
|
||||
|
||||
def named_number(self, name, num=None, default=0):
|
||||
if num is None:
|
||||
return int(self.has_signal(name)) or default
|
||||
else:
|
||||
with open(os.path.join(self.log_dir, name), 'w') as writer:
|
||||
writer.write(str(num))
|
||||
return num
|
||||
|
||||
epoch = functools.partialmethod(named_number, 'epoch')
|
||||
iter = functools.partialmethod(named_number, 'iter')
|
||||
|
||||
def message(self, level, content):
|
||||
self.message_logger.__getattribute__(level)(content)
|
||||
|
||||
def images(self, prefix, image_dict, step):
|
||||
for name, image in image_dict.items():
|
||||
self.add_image(prefix + '/' + name, image, step, dataformats='HWC')
|
||||
|
||||
def merge_save_images(self, name, images):
|
||||
for i, image in enumerate(images):
|
||||
if i == 0:
|
||||
result = image
|
||||
else:
|
||||
result = np.concatenate([result, image], 0)
|
||||
cv2.imwrite(os.path.join(self.vis_dir(), name+'.jpg'), result)
|
||||
|
||||
def vis_dir(self):
|
||||
vis_dir = os.path.join(self.log_dir, self.VISUALIZE_NAME)
|
||||
if not os.path.exists(vis_dir):
|
||||
os.mkdir(vis_dir)
|
||||
return vis_dir
|
||||
|
||||
def save_image_dict(self, images, max_size=1024):
|
||||
for file_name, image in images.items():
|
||||
height, width = image.shape[:2]
|
||||
if height > width:
|
||||
actual_height = min(height, max_size)
|
||||
actual_width = int(round(actual_height * width / height))
|
||||
else:
|
||||
actual_width = min(width, max_size)
|
||||
actual_height = int(round(actual_width * height / width))
|
||||
image = cv2.resize(image, (actual_width, actual_height))
|
||||
cv2.imwrite(os.path.join(self.vis_dir(), file_name+'.jpg'), image)
|
||||
|
||||
def __getattr__(self, name):
|
||||
message_levels = set(['debug', 'info', 'warning', 'error', 'critical'])
|
||||
if name == '__setstate__':
|
||||
raise AttributeError('haha')
|
||||
if name in message_levels:
|
||||
return functools.partial(self.message, name)
|
||||
elif hasattr(self.__dict__.get('tf_board_logger'), name):
|
||||
return self.tf_board_logger.__getattribute__(name)
|
||||
else:
|
||||
super()
|
||||
@@ -0,0 +1,17 @@
|
||||
import os
|
||||
|
||||
|
||||
class SignalMonitor(object):
|
||||
def __init__(self, file_path):
|
||||
self.file_path = file_path
|
||||
|
||||
def get_signal(self):
|
||||
if self.file_path is None:
|
||||
return None
|
||||
if os.path.exists(self.file_path):
|
||||
with open(self.file_path) as f:
|
||||
data = self.file.read()
|
||||
os.remove(f)
|
||||
return data
|
||||
else:
|
||||
return None
|
||||
@@ -0,0 +1,98 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
# File : visualizer.py
|
||||
# Author : Zhaoyi Wan <wanzhaoyi@megvii.com>
|
||||
# Date : 08.01.2019
|
||||
# Last Modified Date: 02.12.2019
|
||||
# Last Modified By : Minghui Liao
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
class Visualize:
|
||||
@classmethod
|
||||
def visualize(cls, x):
|
||||
dimension = len(x.shape)
|
||||
if dimension == 2:
|
||||
pass
|
||||
elif dimension == 3:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def to_np(cls, x):
|
||||
return x.cpu().data.numpy()
|
||||
|
||||
@classmethod
|
||||
def visualize_weights(cls, tensor, format='HW', normalize=True):
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
x = cls.to_np(tensor.permute(format.index('H'), format.index('W')))
|
||||
else:
|
||||
x = tensor.transpose(format.index('H'), format.index('W'))
|
||||
if normalize:
|
||||
x = (x - x.min()) / (x.max() - x.min())
|
||||
# return np.tile(x * 255., (3, 1, 1)).swapaxes(0, 2).swapaxes(1, 0).astype(np.uint8)
|
||||
return cv2.applyColorMap((x * 255).astype(np.uint8), cv2.COLORMAP_JET)
|
||||
|
||||
@classmethod
|
||||
def visualize_points(cls, image, tensor, radius=5, normalized=True):
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
points = cls.to_np(tensor)
|
||||
else:
|
||||
points = tensor
|
||||
if normalized:
|
||||
points = points * image.shape[:2][::-1]
|
||||
for i in range(points.shape[0]):
|
||||
color = np.random.randint(
|
||||
0, 255, (3, ), dtype=np.uint8).astype(np.float)
|
||||
image = cv2.circle(image,
|
||||
tuple(points[i].astype(np.int32).tolist()),
|
||||
radius, color, thickness=radius//2)
|
||||
return image
|
||||
|
||||
@classmethod
|
||||
def visualize_heatmap(cls, tensor, format='CHW'):
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
x = cls.to_np(tensor.permute(format.index('H'),
|
||||
format.index('W'), format.index('C')))
|
||||
else:
|
||||
x = tensor.transpose(
|
||||
format.index('H'), format.index('W'), format.index('C'))
|
||||
canvas = np.zeros((x.shape[0], x.shape[1], 3), dtype=np.float)
|
||||
|
||||
for c in range(0, x.shape[-1]):
|
||||
color = np.random.randint(
|
||||
0, 255, (3, ), dtype=np.uint8).astype(np.float)
|
||||
canvas += np.tile(x[:, :, c], (3, 1, 1)
|
||||
).swapaxes(0, 2).swapaxes(1, 0) * color
|
||||
|
||||
canvas = canvas.astype(np.uint8)
|
||||
return canvas
|
||||
|
||||
@classmethod
|
||||
def visualize_classes(cls, x):
|
||||
canvas = np.zeros((x.shape[0], x.shape[1], 3), dtype=np.uint8)
|
||||
for c in range(int(x.max())):
|
||||
color = np.random.randint(
|
||||
0, 255, (3, ), dtype=np.uint8).astype(np.float)
|
||||
canvas[np.where(x == c)] = color
|
||||
return canvas
|
||||
|
||||
@classmethod
|
||||
def visualize_grid(cls, x, y, stride=16, color=(0, 0, 255), canvas=None):
|
||||
h, w = x.shape
|
||||
if canvas is None:
|
||||
canvas = np.zeros((h, w, 3), dtype=np.uint8)
|
||||
# canvas = np.concatenate([canvas, canvas], axis=1)
|
||||
i, j = 0, 0
|
||||
while i < w:
|
||||
j = 0
|
||||
while j < h:
|
||||
canvas = cv2.circle(canvas, (int(x[i, j] * w + 0.5), int(y[i, j] * h + 0.5)), radius=max(stride//4, 1), color=color, thickness=stride//8)
|
||||
j += stride
|
||||
i += stride
|
||||
return canvas
|
||||
|
||||
@classmethod
|
||||
def visualize_rect(cls, canvas, _rect, color=(0, 0, 255)):
|
||||
rect = (_rect + 0.5).astype(np.int32)
|
||||
return cv2.rectangle(canvas, (rect[0], rect[1]), (rect[2], rect[3]), color)
|
||||
@@ -0,0 +1,19 @@
|
||||
#!/usr/bin/env mdl
|
||||
class WebCV2:
|
||||
def __init__(self):
|
||||
import cv2
|
||||
self._cv2 = cv2
|
||||
from .manager import global_manager as gm
|
||||
self._gm = gm
|
||||
|
||||
def __getattr__(self, name):
|
||||
if hasattr(self._gm, name):
|
||||
return getattr(self._gm, name)
|
||||
elif hasattr(self._cv2, name):
|
||||
return getattr(self._cv2, name)
|
||||
else:
|
||||
raise AttributeError
|
||||
|
||||
import sys
|
||||
sys.modules[__name__] = WebCV2()
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env mdl
|
||||
import socket
|
||||
import base64
|
||||
import cv2
|
||||
import numpy as np
|
||||
from collections import OrderedDict
|
||||
|
||||
from .server import get_server
|
||||
|
||||
|
||||
def jpeg_encode(img):
|
||||
return cv2.imencode('.png', img)[1]
|
||||
|
||||
|
||||
def get_free_port(rng, low=2000, high=10000):
|
||||
in_use = True
|
||||
while in_use:
|
||||
port = rng.randint(high - low) + low
|
||||
in_use = False
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
try:
|
||||
s.bind(("0.0.0.0", port))
|
||||
except socket.error as e:
|
||||
if e.errno == 98: # port already in use
|
||||
in_use = True
|
||||
s.close()
|
||||
return port
|
||||
|
||||
|
||||
class Manager:
|
||||
def __init__(self, img_encode_method=jpeg_encode, rng=None):
|
||||
self._queue = OrderedDict()
|
||||
self._server = None
|
||||
self.img_encode_method = img_encode_method
|
||||
if rng is None:
|
||||
rng = np.random.RandomState(self.get_default_seed())
|
||||
self.rng = rng
|
||||
|
||||
def get_default_seed(self):
|
||||
return 0
|
||||
|
||||
def imshow(self, title, img):
|
||||
data = self.img_encode_method(img)
|
||||
data = base64.b64encode(data)
|
||||
data = data.decode('utf8')
|
||||
self._queue[title] = data
|
||||
|
||||
def waitKey(self, delay=0):
|
||||
if self._server is None:
|
||||
self.port = get_free_port(self.rng)
|
||||
self._server, self._conn = get_server(port=self.port)
|
||||
self._conn.send([delay, list(self._queue.items())])
|
||||
# self._queue = OrderedDict()
|
||||
return self._conn.recv()
|
||||
|
||||
global_manager = Manager()
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env mdl
|
||||
import os
|
||||
BASE_DIR = os.path.dirname(os.path.realpath(__file__))
|
||||
import time
|
||||
import json
|
||||
import select
|
||||
import traceback
|
||||
import socket
|
||||
from multiprocessing import Process, Pipe
|
||||
|
||||
import gevent
|
||||
from gevent.pywsgi import WSGIServer
|
||||
from geventwebsocket.handler import WebSocketHandler
|
||||
from flask import Flask, request, render_template, abort
|
||||
|
||||
|
||||
def log_important_msg(msg, *, padding=3):
|
||||
msg_len = len(msg)
|
||||
width = msg_len + padding * 2 + 2
|
||||
print('#' * width)
|
||||
print('#' + ' ' * (width - 2) + '#')
|
||||
print('#' + ' ' * padding + msg + ' ' * padding + '#')
|
||||
print('#' + ' ' * (width - 2) + '#')
|
||||
print('#' * width)
|
||||
|
||||
|
||||
def hint_url(url, port):
|
||||
log_important_msg(
|
||||
'The server is running at: {}'.format(url))
|
||||
|
||||
|
||||
def _set_server(conn, name='webcv2', port=7788):
|
||||
package = None
|
||||
package_alive = False
|
||||
|
||||
app = Flask(name)
|
||||
app.root_path = BASE_DIR
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
return render_template('index.html', title=name)
|
||||
|
||||
@app.route('/stream')
|
||||
def stream():
|
||||
def poll_ws(ws, delay):
|
||||
return len(select.select([ws.stream.handler.rfile], [], [], delay / 1000.)[0]) > 0
|
||||
|
||||
if request.environ.get('wsgi.websocket'):
|
||||
ws = request.environ['wsgi.websocket']
|
||||
if ws is None:
|
||||
abort(404)
|
||||
else:
|
||||
should_send = True
|
||||
while not ws.closed:
|
||||
global package
|
||||
global package_alive
|
||||
if conn.poll():
|
||||
package = conn.recv()
|
||||
package_alive = True
|
||||
should_send = True
|
||||
if not should_send:
|
||||
continue
|
||||
should_send = False
|
||||
if package is None:
|
||||
ws.send(None)
|
||||
else:
|
||||
delay, info_lst = package
|
||||
ws.send(json.dumps((time.time(), package_alive, delay, info_lst)))
|
||||
if package_alive:
|
||||
if delay <= 0 or poll_ws(ws, delay):
|
||||
message = ws.receive()
|
||||
if ws.closed or message is None:
|
||||
break
|
||||
try:
|
||||
if isinstance(message, bytes):
|
||||
message = message.decode('utf8')
|
||||
message = int(message)
|
||||
except:
|
||||
traceback.print_exc()
|
||||
message = -1
|
||||
else:
|
||||
message = -1
|
||||
conn.send(message)
|
||||
package_alive = False
|
||||
return ""
|
||||
|
||||
http_server = WSGIServer(('', port), app, handler_class=WebSocketHandler)
|
||||
hint_url('http://{}:{}'.format(socket.getfqdn(), port), port)
|
||||
http_server.serve_forever()
|
||||
|
||||
|
||||
def get_server(name='webcv2', port=7788):
|
||||
conn_server, conn_factory = Pipe()
|
||||
p_server = Process(
|
||||
target=_set_server,
|
||||
args=(conn_server,),
|
||||
kwargs=dict(
|
||||
name=name, port=port,
|
||||
),
|
||||
)
|
||||
p_server.daemon = True
|
||||
p_server.start()
|
||||
return p_server, conn_factory
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta http-equiv="Content-Type" content="text/html; charset=UTF-8"/>
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1, maximum-scale=1.0, user-scalable=no"/>
|
||||
<title>{{title}}</title>
|
||||
<link rel="icon" href="//assets.megvii.com/static%2Ffavicon.ico?ver=1498037959257">
|
||||
<script src="//cdnjs.cloudflare.com/ajax/libs/jquery/3.2.1/jquery.min.js"></script>
|
||||
<script src="//cdnjs.cloudflare.com/ajax/libs/vue/2.2.6/vue.min.js"></script>
|
||||
</head>
|
||||
|
||||
<body>
|
||||
<script language="javascript" type="text/javascript">
|
||||
var vm = null;
|
||||
function isCharacterKeyPress(evt) {
|
||||
if (typeof evt.which == "undefined") {
|
||||
// This is IE, which only fires keypress events for printable keys
|
||||
return true;
|
||||
} else if (typeof evt.which == "number" && evt.which > 0) {
|
||||
// In other browsers except old versions of WebKit, evt.which is
|
||||
// only greater than zero if the keypress is a printable key.
|
||||
// We need to filter out backspace and ctrl/alt/meta key combinations
|
||||
return !evt.ctrlKey && !evt.metaKey && !evt.altKey && evt.which != 8;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
$(function() {
|
||||
vm = new Vue({
|
||||
el: '#board',
|
||||
data: {
|
||||
websocket: null,
|
||||
// resources
|
||||
delay: -1,
|
||||
data_alive: null,
|
||||
img_lst: null,
|
||||
// network
|
||||
package_size: 0,
|
||||
download_time: 0,
|
||||
net_speed: 0,
|
||||
// timer
|
||||
timer_interval: null,
|
||||
timer_step: 100,
|
||||
},
|
||||
created: function() {
|
||||
this.create_connection();
|
||||
document.addEventListener("keypress", this.send_key, false);
|
||||
},
|
||||
computed: {
|
||||
connection_alive: function() {
|
||||
return (this.websocket != null) &&
|
||||
!(this.websocket.readyState == this.websocket.CLOSING || this.websocket.readyState == this.websocket.CLOSED);
|
||||
},
|
||||
},
|
||||
methods: {
|
||||
get_ws_url: function(s) {
|
||||
var l = window.location;
|
||||
return ((l.protocol === "https:") ? "wss://" : "ws://") + l.hostname + (((l.port != 80) && (l.port != 443)) ? ":" + l.port : "") + l.pathname + s;
|
||||
},
|
||||
create_connection: function() {
|
||||
this.websocket = new WebSocket(this.get_ws_url("stream"));
|
||||
this.websocket.binaryType = "arraybuffer";
|
||||
this.websocket.onopen = function(evt) { console.log("connected"); };
|
||||
this.websocket.onclose = this.close_connection;
|
||||
this.websocket.onerror = function(evt) { console.log("error occurred"); console.log(evt); };
|
||||
this.websocket.onmessage = this.receive_message;
|
||||
},
|
||||
close_connection: function() {
|
||||
console.log("disconnected");
|
||||
this.websocket = null;
|
||||
},
|
||||
receive_message: function(evt) {
|
||||
var obj = JSON.parse(evt.data);
|
||||
var send_time = obj[0], recv_time = (new Date()).getTime() / 1000;
|
||||
this.package_size = evt.data.length * 8;
|
||||
this.download_time = recv_time - send_time;
|
||||
this.net_speed = this.package_size / this.download_time;
|
||||
console.log(obj);
|
||||
this.data_alive = obj[1];
|
||||
this.delay = obj[2];
|
||||
this.img_lst = obj[3];
|
||||
this.start_timer();
|
||||
},
|
||||
send_key: function(evt) {
|
||||
if (!this.connection_alive || !this.data_alive) return;
|
||||
if (!isCharacterKeyPress(evt)) return;
|
||||
this.delay = 0;
|
||||
this.data_alive = false;
|
||||
var keycode = evt.keyCode;
|
||||
console.log("key pressed " + keycode + " " + String.fromCharCode(keycode));
|
||||
this.websocket.send(keycode);
|
||||
},
|
||||
update_timer: function() {
|
||||
if (this.delay > 0) this.delay -= this.timer_step;
|
||||
if (this.delay <= 0) {
|
||||
this.data_alive = false;
|
||||
if (this.timer_interval != null) {
|
||||
clearInterval(this.timer_interval);
|
||||
this.timer_interval = null;
|
||||
}
|
||||
}
|
||||
},
|
||||
start_timer: function() {
|
||||
if (this.timer_interval != null) {
|
||||
clearInterval(this.timer_interval);
|
||||
this.timer_interval = null;
|
||||
}
|
||||
if (this.delay > 0) this.timer_interval = setInterval(this.update_timer, this.timer_step);
|
||||
},
|
||||
},
|
||||
filters: {
|
||||
decimal: function(num, fixed_point) {
|
||||
return num.toFixed(fixed_point);
|
||||
},
|
||||
},
|
||||
});
|
||||
});
|
||||
</script>
|
||||
|
||||
{% raw %}
|
||||
<!-- Vue template -->
|
||||
<div id="board">
|
||||
<template v-if="connection_alive">
|
||||
<template v-if="data_alive">
|
||||
<p v-if="delay > 0">Press any key in {{delay / 1000 | decimal(1)}} seconds. </p>
|
||||
<p v-else>Press any key to continue. </p>
|
||||
</template>
|
||||
<template v-else>
|
||||
<p>Waiting for response from server. </p>
|
||||
</template>
|
||||
</template>
|
||||
<template v-else>
|
||||
<p>Disconnected from server. </p>
|
||||
</template>
|
||||
<p>Network: {{net_speed / 1000000 | decimal(2)}} MB/s * {{download_time | decimal(2)}} s = {{package_size / 1000000 | decimal(2)}} MB</p>
|
||||
<div style="display:inline-block">
|
||||
<div style="display:inline-block; vertical-align: top;" v-for="obj in img_lst">
|
||||
<p>{{obj[0]}}</p>
|
||||
<img :src="'data:image/png;base64,' + obj[1]" />
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endraw %}
|
||||
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,32 @@
|
||||
from detectron2.config import CfgNode as CN
|
||||
|
||||
|
||||
def add_vit_config(cfg):
|
||||
"""
|
||||
Add config for VIT.
|
||||
"""
|
||||
_C = cfg
|
||||
|
||||
_C.MODEL.VIT = CN()
|
||||
|
||||
# CoaT model name.
|
||||
_C.MODEL.VIT.NAME = ""
|
||||
|
||||
# Output features from CoaT backbone.
|
||||
_C.MODEL.VIT.OUT_FEATURES = ["layer3", "layer5", "layer7", "layer11"]
|
||||
|
||||
_C.MODEL.VIT.IMG_SIZE = [224, 224]
|
||||
|
||||
_C.MODEL.VIT.POS_TYPE = "shared_rel"
|
||||
|
||||
_C.MODEL.VIT.DROP_PATH = 0.
|
||||
|
||||
_C.MODEL.VIT.MODEL_KWARGS = "{}"
|
||||
|
||||
_C.SOLVER.OPTIMIZER = "ADAMW"
|
||||
|
||||
_C.SOLVER.BACKBONE_MULTIPLIER = 1.0
|
||||
|
||||
_C.AUG = CN()
|
||||
|
||||
_C.AUG.DETR = False
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
# from https://github.com/facebookresearch/detr/blob/main/d2/detr/dataset_mapper.py
|
||||
|
||||
|
||||
import copy
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from detectron2.data import detection_utils as utils
|
||||
from detectron2.data import transforms as T
|
||||
|
||||
__all__ = ["DetrDatasetMapper"]
|
||||
|
||||
|
||||
def build_transform_gen(cfg, is_train):
|
||||
"""
|
||||
Create a list of :class:`TransformGen` from config.
|
||||
Returns:
|
||||
list[TransformGen]
|
||||
"""
|
||||
if is_train:
|
||||
min_size = cfg.INPUT.MIN_SIZE_TRAIN
|
||||
max_size = cfg.INPUT.MAX_SIZE_TRAIN
|
||||
sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING
|
||||
else:
|
||||
min_size = cfg.INPUT.MIN_SIZE_TEST
|
||||
max_size = cfg.INPUT.MAX_SIZE_TEST
|
||||
sample_style = "choice"
|
||||
if sample_style == "range":
|
||||
assert len(min_size) == 2, "more than 2 ({}) min_size(s) are provided for ranges".format(len(min_size))
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
tfm_gens = []
|
||||
if is_train:
|
||||
tfm_gens.append(T.RandomFlip())
|
||||
tfm_gens.append(T.ResizeShortestEdge(min_size, max_size, sample_style))
|
||||
if is_train:
|
||||
logger.info("TransformGens used in training: " + str(tfm_gens))
|
||||
return tfm_gens
|
||||
|
||||
|
||||
class DetrDatasetMapper:
|
||||
"""
|
||||
A callable which takes a dataset dict in Detectron2 Dataset format,
|
||||
and map it into a format used by DETR.
|
||||
|
||||
The callable currently does the following:
|
||||
|
||||
1. Read the image from "file_name"
|
||||
2. Applies geometric transforms to the image and annotation
|
||||
3. Find and applies suitable cropping to the image and annotation
|
||||
4. Prepare image and annotation to Tensors
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, is_train=True):
|
||||
if cfg.INPUT.CROP.ENABLED and is_train:
|
||||
self.crop_gen = [
|
||||
T.ResizeShortestEdge([400, 500, 600], sample_style="choice"),
|
||||
T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE),
|
||||
]
|
||||
else:
|
||||
self.crop_gen = None
|
||||
|
||||
self.mask_on = cfg.MODEL.MASK_ON
|
||||
self.tfm_gens = build_transform_gen(cfg, is_train)
|
||||
logging.getLogger(__name__).info(
|
||||
"Full TransformGens used in training: {}, crop: {}".format(str(self.tfm_gens), str(self.crop_gen))
|
||||
)
|
||||
|
||||
self.img_format = cfg.INPUT.FORMAT
|
||||
self.is_train = is_train
|
||||
|
||||
def __call__(self, dataset_dict):
|
||||
"""
|
||||
Args:
|
||||
dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
|
||||
|
||||
Returns:
|
||||
dict: a format that builtin models in detectron2 accept
|
||||
"""
|
||||
dataset_dict = copy.deepcopy(dataset_dict) # it will be modified by code below
|
||||
image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
|
||||
utils.check_image_size(dataset_dict, image)
|
||||
|
||||
if self.crop_gen is None:
|
||||
image, transforms = T.apply_transform_gens(self.tfm_gens, image)
|
||||
else:
|
||||
if np.random.rand() > 0.5:
|
||||
image, transforms = T.apply_transform_gens(self.tfm_gens, image)
|
||||
else:
|
||||
image, transforms = T.apply_transform_gens(
|
||||
self.tfm_gens[:-1] + self.crop_gen + self.tfm_gens[-1:], image
|
||||
)
|
||||
|
||||
image_shape = image.shape[:2] # h, w
|
||||
|
||||
# Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,
|
||||
# but not efficient on large generic data structures due to the use of pickle & mp.Queue.
|
||||
# Therefore it's important to use torch.Tensor.
|
||||
dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
|
||||
|
||||
if not self.is_train:
|
||||
# USER: Modify this if you want to keep them for some reason.
|
||||
dataset_dict.pop("annotations", None)
|
||||
return dataset_dict
|
||||
|
||||
if "annotations" in dataset_dict:
|
||||
# USER: Modify this if you want to keep them for some reason.
|
||||
for anno in dataset_dict["annotations"]:
|
||||
if not self.mask_on:
|
||||
anno.pop("segmentation", None)
|
||||
anno.pop("keypoints", None)
|
||||
|
||||
# USER: Implement additional transformations if you have other types of data
|
||||
annos = [
|
||||
utils.transform_instance_annotations(obj, transforms, image_shape)
|
||||
for obj in dataset_dict.pop("annotations")
|
||||
if obj.get("iscrowd", 0) == 0
|
||||
]
|
||||
instances = utils.annotations_to_instances(annos, image_shape)
|
||||
dataset_dict["instances"] = utils.filter_empty_instances(instances)
|
||||
return dataset_dict
|
||||
@@ -0,0 +1,476 @@
|
||||
"""
|
||||
Mostly copy-paste from DINO and timm library:
|
||||
https://github.com/facebookresearch/dino
|
||||
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
|
||||
"""
|
||||
import warnings
|
||||
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint as checkpoint
|
||||
from timm.models.layers import trunc_normal_, drop_path, to_2tuple
|
||||
from functools import partial
|
||||
|
||||
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 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.):
|
||||
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
|
||||
q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads,
|
||||
C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
|
||||
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, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
|
||||
self.window_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
|
||||
self.num_patches_w, self.num_patches_h = self.window_size
|
||||
|
||||
self.num_patches = self.window_size[0] * self.window_size[1]
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim,
|
||||
kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.proj(x)
|
||||
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 ViT(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
model_name='vit_base_patch16_224',
|
||||
img_size=384,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
num_classes=19,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
qk_scale=None,
|
||||
drop_rate=0.1,
|
||||
attn_drop_rate=0.,
|
||||
drop_path_rate=0.,
|
||||
hybrid_backbone=None,
|
||||
norm_layer=partial(nn.LayerNorm, eps=1e-6),
|
||||
norm_cfg=None,
|
||||
pos_embed_interp=False,
|
||||
random_init=False,
|
||||
align_corners=False,
|
||||
use_checkpoint=False,
|
||||
num_extra_tokens=1,
|
||||
out_features=None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
super(ViT, self).__init__()
|
||||
self.model_name = model_name
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.in_chans = in_chans
|
||||
self.embed_dim = embed_dim
|
||||
self.depth = depth
|
||||
self.num_heads = num_heads
|
||||
self.num_classes = num_classes
|
||||
self.mlp_ratio = mlp_ratio
|
||||
self.qkv_bias = qkv_bias
|
||||
self.qk_scale = qk_scale
|
||||
self.drop_rate = drop_rate
|
||||
self.attn_drop_rate = attn_drop_rate
|
||||
self.drop_path_rate = drop_path_rate
|
||||
self.hybrid_backbone = hybrid_backbone
|
||||
self.norm_layer = norm_layer
|
||||
self.norm_cfg = norm_cfg
|
||||
self.pos_embed_interp = pos_embed_interp
|
||||
self.random_init = random_init
|
||||
self.align_corners = align_corners
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.num_extra_tokens = num_extra_tokens
|
||||
self.out_features = out_features
|
||||
self.out_indices = [int(name[5:]) for name in out_features]
|
||||
|
||||
# self.num_stages = self.depth
|
||||
# self.out_indices = tuple(range(self.num_stages))
|
||||
|
||||
if self.hybrid_backbone is not None:
|
||||
self.patch_embed = HybridEmbed(
|
||||
self.hybrid_backbone, img_size=self.img_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
|
||||
else:
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=self.img_size, patch_size=self.patch_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
|
||||
self.num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
|
||||
|
||||
if self.num_extra_tokens == 2:
|
||||
self.dist_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
|
||||
|
||||
self.pos_embed = nn.Parameter(torch.zeros(
|
||||
1, self.num_patches + self.num_extra_tokens, self.embed_dim))
|
||||
self.pos_drop = nn.Dropout(p=self.drop_rate)
|
||||
|
||||
# self.num_extra_tokens = self.pos_embed.shape[-2] - self.num_patches
|
||||
dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate,
|
||||
self.depth)] # stochastic depth decay rule
|
||||
self.blocks = nn.ModuleList([
|
||||
Block(
|
||||
dim=self.embed_dim, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, qkv_bias=self.qkv_bias,
|
||||
qk_scale=self.qk_scale,
|
||||
drop=self.drop_rate, attn_drop=self.attn_drop_rate, drop_path=dpr[i], norm_layer=self.norm_layer)
|
||||
for i in range(self.depth)])
|
||||
|
||||
# NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
|
||||
# self.repr = nn.Linear(embed_dim, representation_size)
|
||||
# self.repr_act = nn.Tanh()
|
||||
|
||||
if patch_size == 16:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
nn.SyncBatchNorm(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn3 = nn.Identity()
|
||||
|
||||
self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
elif patch_size == 8:
|
||||
self.fpn1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn2 = nn.Identity()
|
||||
|
||||
self.fpn3 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||||
)
|
||||
|
||||
self.fpn4 = nn.Sequential(
|
||||
nn.MaxPool2d(kernel_size=4, stride=4),
|
||||
)
|
||||
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
if self.num_extra_tokens==2:
|
||||
trunc_normal_(self.dist_token, std=0.2)
|
||||
self.apply(self._init_weights)
|
||||
# self.fix_init_weight()
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
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 init_weights(self):
|
||||
logger = get_root_logger()
|
||||
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
if self.init_cfg is None:
|
||||
logger.warn(f'No pre-trained weights for '
|
||||
f'{self.__class__.__name__}, '
|
||||
f'training start from scratch')
|
||||
else:
|
||||
assert 'checkpoint' in self.init_cfg, f'Only support ' \
|
||||
f'specify `Pretrained` in ' \
|
||||
f'`init_cfg` in ' \
|
||||
f'{self.__class__.__name__} '
|
||||
logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
|
||||
load_checkpoint(self, filename=self.init_cfg['checkpoint'], strict=False, logger=logger)
|
||||
'''
|
||||
|
||||
def get_num_layers(self):
|
||||
return len(self.blocks)
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'pos_embed', 'cls_token'}
|
||||
|
||||
def _conv_filter(self, 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
|
||||
|
||||
def to_2D(self, x):
|
||||
n, hw, c = x.shape
|
||||
h = w = int(math.sqrt(hw))
|
||||
x = x.transpose(1, 2).reshape(n, c, h, w)
|
||||
return x
|
||||
|
||||
def to_1D(self, x):
|
||||
n, c, h, w = x.shape
|
||||
x = x.reshape(n, c, -1).transpose(1, 2)
|
||||
return x
|
||||
|
||||
def interpolate_pos_encoding(self, x, w, h):
|
||||
npatch = x.shape[1] - self.num_extra_tokens
|
||||
N = self.pos_embed.shape[1] - self.num_extra_tokens
|
||||
if npatch == N and w == h:
|
||||
return self.pos_embed
|
||||
|
||||
class_ORdist_pos_embed = self.pos_embed[:, 0:self.num_extra_tokens]
|
||||
|
||||
patch_pos_embed = self.pos_embed[:, self.num_extra_tokens:]
|
||||
|
||||
dim = x.shape[-1]
|
||||
w0 = w // self.patch_embed.patch_size[0]
|
||||
h0 = h // self.patch_embed.patch_size[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 = w0 + 0.1, h0 + 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=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
|
||||
mode='bicubic',
|
||||
)
|
||||
assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
|
||||
return torch.cat((class_ORdist_pos_embed, patch_pos_embed), dim=1)
|
||||
|
||||
def prepare_tokens(self, x, mask=None):
|
||||
B, nc, w, h = x.shape
|
||||
# patch linear embedding
|
||||
x = self.patch_embed(x)
|
||||
|
||||
# mask image modeling
|
||||
if mask is not None:
|
||||
x = self.mask_model(x, mask)
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
|
||||
# add the [CLS] token to the embed patch tokens
|
||||
all_tokens = [self.cls_token.expand(B, -1, -1)]
|
||||
|
||||
if self.num_extra_tokens == 2:
|
||||
dist_tokens = self.dist_token.expand(B, -1, -1)
|
||||
all_tokens.append(dist_tokens)
|
||||
all_tokens.append(x)
|
||||
|
||||
x = torch.cat(all_tokens, dim=1)
|
||||
|
||||
# add positional encoding to each token
|
||||
x = x + self.interpolate_pos_encoding(x, w, h)
|
||||
|
||||
return self.pos_drop(x)
|
||||
|
||||
def forward_features(self, x):
|
||||
# print(f"==========shape of x is {x.shape}==========")
|
||||
B, _, H, W = x.shape
|
||||
Hp, Wp = H // self.patch_size, W // self.patch_size
|
||||
x = self.prepare_tokens(x)
|
||||
|
||||
features = []
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if self.use_checkpoint:
|
||||
x = checkpoint.checkpoint(blk, x)
|
||||
else:
|
||||
x = blk(x)
|
||||
if i in self.out_indices:
|
||||
xp = x[:, self.num_extra_tokens:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
|
||||
features.append(xp.contiguous())
|
||||
|
||||
ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
|
||||
for i in range(len(features)):
|
||||
features[i] = ops[i](features[i])
|
||||
|
||||
feat_out = {}
|
||||
|
||||
for name, value in zip(self.out_features, features):
|
||||
feat_out[name] = value
|
||||
|
||||
return feat_out
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
return x
|
||||
|
||||
|
||||
def deit_base_patch16(pretrained=False, **kwargs):
|
||||
model = ViT(
|
||||
patch_size=16,
|
||||
drop_rate=0.,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
num_classes=1000,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
use_checkpoint=True,
|
||||
num_extra_tokens=2,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
|
||||
def mae_base_patch16(pretrained=False, **kwargs):
|
||||
model = ViT(
|
||||
patch_size=16,
|
||||
drop_rate=0.,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
num_classes=1000,
|
||||
mlp_ratio=4.,
|
||||
qkv_bias=True,
|
||||
use_checkpoint=True,
|
||||
num_extra_tokens=1,
|
||||
**kwargs)
|
||||
model.default_cfg = _cfg()
|
||||
return model
|
||||
@@ -0,0 +1,128 @@
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
import itertools
|
||||
from collections import OrderedDict
|
||||
|
||||
import detectron2.utils.comm as comm
|
||||
from detectron2.evaluation import COCOEvaluator
|
||||
|
||||
from .concern.icdar2015_eval.detection.iou import DetectionIoUEvaluator
|
||||
|
||||
class FUNSDEvaluator(COCOEvaluator):
|
||||
def evaluate(self, img_ids=None):
|
||||
"""
|
||||
Args:
|
||||
img_ids: a list of image IDs to evaluate on. Default to None for the whole dataset
|
||||
"""
|
||||
if self._distributed:
|
||||
comm.synchronize()
|
||||
predictions = comm.gather(self._predictions, dst=0)
|
||||
predictions = list(itertools.chain(*predictions))
|
||||
|
||||
if not comm.is_main_process():
|
||||
return {}
|
||||
else:
|
||||
predictions = self._predictions
|
||||
|
||||
if len(predictions) == 0:
|
||||
self._logger.warning("[COCOEvaluator] Did not receive valid predictions.")
|
||||
return {}
|
||||
|
||||
self._logger.warning("[evaluating...]The evaluator may take long time")
|
||||
|
||||
id2img = {}
|
||||
gt = {}
|
||||
with open('data/instances_test.json', 'r',
|
||||
encoding='utf-8') as fr:
|
||||
data = json.load(fr)
|
||||
for img in data['images']:
|
||||
id = img['id']
|
||||
name = os.path.basename(img['file_name'])[:-len('.jpg')]
|
||||
assert id not in id2img.keys()
|
||||
id2img[id] = name
|
||||
assert len(id2img) == len(data['images'])
|
||||
|
||||
img2id, id2bbox = {}, {}
|
||||
for i in range(len(data['images'])):
|
||||
key = os.path.basename(data['images'][i]['file_name'][:-len('.png')])
|
||||
assert key not in img2id.keys()
|
||||
img2id[key] = data['images'][i]['id']
|
||||
for i in range(len(data['annotations'])):
|
||||
img_id = data['annotations'][i]['image_id']
|
||||
if img_id not in id2bbox.keys():
|
||||
id2bbox[img_id] = []
|
||||
x0, y0, w, h = data['annotations'][i]['bbox']
|
||||
x1, y1 = x0 + w, y0 + h
|
||||
line = [(x0, y0), (x1, y0), (x1, y1), (x0, y1)]
|
||||
id2bbox[img_id].append(
|
||||
{
|
||||
'points': line,
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}
|
||||
)
|
||||
for key, val in img2id.items():
|
||||
assert key not in gt.keys()
|
||||
gt[key] = id2bbox[val]
|
||||
|
||||
self._results = OrderedDict()
|
||||
|
||||
evaluator = DetectionIoUEvaluator()
|
||||
|
||||
for iter in range(3, 10):
|
||||
thr = iter * 0.1
|
||||
self._results[thr] = {}
|
||||
|
||||
total_prediction = {}
|
||||
for cur_pred in predictions:
|
||||
assert cur_pred['image_id'] in id2img.keys()
|
||||
id = id2img[cur_pred['image_id']]
|
||||
if id not in total_prediction.keys(): total_prediction[id] = []
|
||||
|
||||
for cur_inst in cur_pred['instances']:
|
||||
x0, y0, w, h = cur_inst['bbox']
|
||||
cur_score = cur_inst['score']
|
||||
if cur_score < thr:
|
||||
continue
|
||||
|
||||
x1, y1 = x0 + w, y0 + h
|
||||
|
||||
x0, x1 = int(x0 + 0.5), int(x1 + 0.5)
|
||||
y0, y1 = int(y0 + 0.5), int(y1 + 0.5)
|
||||
|
||||
min_x, max_x = min([x0, x1]), max([x0, x1])
|
||||
min_y, max_y = min([y0, y1]), max([y0, y1])
|
||||
|
||||
pred_line = [min_x, min_y, max_x, min_y, max_x, max_y, min_x, max_y]
|
||||
pred_line_str = ','.join(list(map(str, pred_line)))
|
||||
|
||||
total_prediction[id].append(pred_line_str)
|
||||
|
||||
final_gt = []
|
||||
final_res = []
|
||||
for key, _ in gt.items():
|
||||
final_gt.append(copy.deepcopy(gt[key]))
|
||||
|
||||
cur_res = []
|
||||
pred = total_prediction[key]
|
||||
for i in range(len(pred)):
|
||||
line = list(map(int, pred[i].split(',')))
|
||||
line = [(line[0], line[1]), (line[2], line[3]), (line[4], line[5]), (line[6], line[7])]
|
||||
cur_res.append(
|
||||
{
|
||||
'points': line,
|
||||
'text': 1234,
|
||||
'ignore': False,
|
||||
}
|
||||
)
|
||||
final_res.append(cur_res)
|
||||
|
||||
results = []
|
||||
for cur_gt, pred in zip(final_gt, final_res):
|
||||
results.append(evaluator.evaluate_image(cur_gt, pred))
|
||||
metrics = evaluator.combine_results(results)
|
||||
for key, val in metrics.items():
|
||||
self._results["{:.1f}_{}".format(thr, key)] = val
|
||||
|
||||
return copy.deepcopy(self._results)
|
||||
@@ -0,0 +1,259 @@
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
|
||||
from typing import Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
#from fvcore.common.checkpoint import _IncompatibleKeys, _strip_prefix_if_present, TORCH_VERSION, quantization, \
|
||||
# ObserverBase, FakeQuantizeBase
|
||||
from fvcore.common.checkpoint import _IncompatibleKeys, _strip_prefix_if_present, TORCH_VERSION
|
||||
from torch import distributed as dist
|
||||
from scipy import interpolate
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def append_prefix(k):
|
||||
prefix = 'backbone.bottom_up.backbone.'
|
||||
# return prefix + k if not k.startswith(prefix) else k
|
||||
return k
|
||||
|
||||
|
||||
def modify_ckpt_state(model, state_dict, logger=None):
|
||||
# reshape absolute position embedding for Swin
|
||||
if state_dict.get(append_prefix('absolute_pos_embed')) is not None:
|
||||
absolute_pos_embed = state_dict[append_prefix('absolute_pos_embed')]
|
||||
N1, L, C1 = absolute_pos_embed.size()
|
||||
N2, C2, H, W = model.backbone.bottom_up.backbone.absolute_pos_embed.size()
|
||||
if N1 != N2 or C1 != C2 or L != H * W:
|
||||
logger.warning("Error in loading absolute_pos_embed, pass")
|
||||
else:
|
||||
state_dict[append_prefix('absolute_pos_embed')] = absolute_pos_embed.view(N2, H, W, C2).permute(0, 3, 1, 2)
|
||||
|
||||
def get_dist_info():
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
rank, _ = get_dist_info()
|
||||
all_keys = list(state_dict.keys())
|
||||
for key in all_keys:
|
||||
if "relative_position_index" in key:
|
||||
state_dict.pop(key)
|
||||
|
||||
if "relative_position_bias_table" in key:
|
||||
rel_pos_bias = state_dict[key]
|
||||
src_num_pos, num_attn_heads = rel_pos_bias.size()
|
||||
if key not in model.state_dict():
|
||||
continue
|
||||
dst_num_pos, _ = model.state_dict()[key].size()
|
||||
dst_patch_shape = model.backbone.bottom_up.backbone.patch_embed.patch_shape
|
||||
if dst_patch_shape[0] != dst_patch_shape[1]:
|
||||
raise NotImplementedError()
|
||||
num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
|
||||
src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
|
||||
dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
|
||||
if src_size != dst_size:
|
||||
if rank == 0:
|
||||
print("Position interpolate for %s from %dx%d to %dx%d" % (
|
||||
key, src_size, src_size, dst_size, dst_size))
|
||||
extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
|
||||
rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
|
||||
|
||||
def geometric_progression(a, r, n):
|
||||
return a * (1.0 - r ** n) / (1.0 - r)
|
||||
|
||||
left, right = 1.01, 1.5
|
||||
while right - left > 1e-6:
|
||||
q = (left + right) / 2.0
|
||||
gp = geometric_progression(1, q, src_size // 2)
|
||||
if gp > dst_size // 2:
|
||||
right = q
|
||||
else:
|
||||
left = q
|
||||
|
||||
# if q > 1.13492:
|
||||
# q = 1.13492
|
||||
|
||||
dis = []
|
||||
cur = 1
|
||||
for i in range(src_size // 2):
|
||||
dis.append(cur)
|
||||
cur += q ** (i + 1)
|
||||
|
||||
r_ids = [-_ for _ in reversed(dis)]
|
||||
|
||||
x = r_ids + [0] + dis
|
||||
y = r_ids + [0] + dis
|
||||
|
||||
t = dst_size // 2.0
|
||||
dx = np.arange(-t, t + 0.1, 1.0)
|
||||
dy = np.arange(-t, t + 0.1, 1.0)
|
||||
if rank == 0:
|
||||
print("x = {}".format(x))
|
||||
print("dx = {}".format(dx))
|
||||
|
||||
all_rel_pos_bias = []
|
||||
|
||||
for i in range(num_attn_heads):
|
||||
z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
|
||||
f = interpolate.interp2d(x, y, z, kind='cubic')
|
||||
all_rel_pos_bias.append(
|
||||
torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
|
||||
|
||||
rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
|
||||
new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
|
||||
state_dict[key] = new_rel_pos_bias
|
||||
|
||||
if append_prefix('pos_embed') in state_dict:
|
||||
pos_embed_checkpoint = state_dict[append_prefix('pos_embed')]
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.backbone.bottom_up.backbone.patch_embed.num_patches
|
||||
num_extra_tokens = model.backbone.bottom_up.backbone.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
# new_size = int(num_patches ** 0.5)
|
||||
new_size_w = model.backbone.bottom_up.backbone.patch_embed.num_patches_w
|
||||
new_size_h = model.backbone.bottom_up.backbone.patch_embed.num_patches_h
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size_h or orig_size != new_size_w:
|
||||
if rank == 0:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size_w, new_size_h))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size_w, new_size_h), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
state_dict[append_prefix('pos_embed')] = new_pos_embed
|
||||
|
||||
# interpolate position bias table if needed
|
||||
relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
|
||||
for table_key in relative_position_bias_table_keys:
|
||||
table_pretrained = state_dict[table_key]
|
||||
if table_key not in model.state_dict():
|
||||
continue
|
||||
table_current = model.state_dict()[table_key]
|
||||
L1, nH1 = table_pretrained.size()
|
||||
L2, nH2 = table_current.size()
|
||||
if nH1 != nH2:
|
||||
logger.warning(f"Error in loading {table_key}, pass")
|
||||
else:
|
||||
if L1 != L2:
|
||||
S1 = int(L1 ** 0.5)
|
||||
S2 = int(L2 ** 0.5)
|
||||
table_pretrained_resized = F.interpolate(
|
||||
table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
|
||||
size=(S2, S2), mode='bicubic')
|
||||
state_dict[table_key] = table_pretrained_resized.view(nH2, L2).permute(1, 0)
|
||||
|
||||
if append_prefix('rel_pos_bias.relative_position_bias_table') in state_dict and \
|
||||
model.backbone.bottom_up.backbone.use_rel_pos_bias and \
|
||||
not model.backbone.bottom_up.backbone.use_shared_rel_pos_bias and \
|
||||
append_prefix('blocks.0.attn.relative_position_bias_table') not in state_dict:
|
||||
logger.info("[BEIT] Expand the shared relative position embedding to each transformer block. ")
|
||||
num_layers = model.backbone.bottom_up.backbone.get_num_layers()
|
||||
rel_pos_bias = state_dict[append_prefix("rel_pos_bias.relative_position_bias_table")]
|
||||
for i in range(num_layers):
|
||||
state_dict["blocks.%d.attn.relative_position_bias_table" % i] = rel_pos_bias.clone()
|
||||
state_dict.pop(append_prefix("rel_pos_bias.relative_position_bias_table"))
|
||||
|
||||
return state_dict
|
||||
|
||||
|
||||
class MyDetectionCheckpointer(DetectionCheckpointer):
|
||||
def _load_model(self, checkpoint: Any) -> _IncompatibleKeys:
|
||||
"""
|
||||
Load weights from a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint (Any): checkpoint contains the weights.
|
||||
|
||||
Returns:
|
||||
``NamedTuple`` with ``missing_keys``, ``unexpected_keys``,
|
||||
and ``incorrect_shapes`` fields:
|
||||
* **missing_keys** is a list of str containing the missing keys
|
||||
* **unexpected_keys** is a list of str containing the unexpected keys
|
||||
* **incorrect_shapes** is a list of (key, shape in checkpoint, shape in model)
|
||||
|
||||
This is just like the return value of
|
||||
:func:`torch.nn.Module.load_state_dict`, but with extra support
|
||||
for ``incorrect_shapes``.
|
||||
"""
|
||||
checkpoint_state_dict = checkpoint.pop("model")
|
||||
self._convert_ndarray_to_tensor(checkpoint_state_dict)
|
||||
|
||||
# if the state_dict comes from a model that was wrapped in a
|
||||
# DataParallel or DistributedDataParallel during serialization,
|
||||
# remove the "module" prefix before performing the matching.
|
||||
_strip_prefix_if_present(checkpoint_state_dict, "module.")
|
||||
|
||||
# workaround https://github.com/pytorch/pytorch/issues/24139
|
||||
model_state_dict = self.model.state_dict()
|
||||
incorrect_shapes = []
|
||||
|
||||
# rename the para in checkpoint_state_dict
|
||||
# some bug here, do not support re load
|
||||
|
||||
checkpoint_state_dict = {
|
||||
append_prefix(k): checkpoint_state_dict[k]
|
||||
for k in checkpoint_state_dict.keys()
|
||||
}
|
||||
|
||||
checkpoint_state_dict = modify_ckpt_state(self.model, checkpoint_state_dict, logger=self.logger)
|
||||
|
||||
for k in list(checkpoint_state_dict.keys()):
|
||||
if k in model_state_dict:
|
||||
model_param = model_state_dict[k]
|
||||
# Allow mismatch for uninitialized parameters
|
||||
if TORCH_VERSION >= (1, 8) and isinstance(
|
||||
model_param, nn.parameter.UninitializedParameter
|
||||
):
|
||||
continue
|
||||
shape_model = tuple(model_param.shape)
|
||||
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
|
||||
if shape_model != shape_checkpoint:
|
||||
|
||||
has_observer_base_classes = (
|
||||
TORCH_VERSION >= (1, 8)
|
||||
and hasattr(quantization, "ObserverBase")
|
||||
and hasattr(quantization, "FakeQuantizeBase")
|
||||
)
|
||||
if has_observer_base_classes:
|
||||
# Handle the special case of quantization per channel observers,
|
||||
# where buffer shape mismatches are expected.
|
||||
def _get_module_for_key(
|
||||
model: torch.nn.Module, key: str
|
||||
) -> torch.nn.Module:
|
||||
# foo.bar.param_or_buffer_name -> [foo, bar]
|
||||
key_parts = key.split(".")[:-1]
|
||||
cur_module = model
|
||||
for key_part in key_parts:
|
||||
cur_module = getattr(cur_module, key_part)
|
||||
return cur_module
|
||||
|
||||
cls_to_skip = (
|
||||
ObserverBase,
|
||||
FakeQuantizeBase,
|
||||
)
|
||||
target_module = _get_module_for_key(self.model, k)
|
||||
if isinstance(target_module, cls_to_skip):
|
||||
# Do not remove modules with expected shape mismatches
|
||||
# them from the state_dict loading. They have special logic
|
||||
# in _load_from_state_dict to handle the mismatches.
|
||||
continue
|
||||
|
||||
incorrect_shapes.append((k, shape_checkpoint, shape_model))
|
||||
checkpoint_state_dict.pop(k)
|
||||
incompatible = self.model.load_state_dict(checkpoint_state_dict, strict=False)
|
||||
return _IncompatibleKeys(
|
||||
missing_keys=incompatible.missing_keys,
|
||||
unexpected_keys=incompatible.unexpected_keys,
|
||||
incorrect_shapes=incorrect_shapes,
|
||||
)
|
||||
@@ -0,0 +1,740 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
|
||||
"""
|
||||
This file contains components with some default boilerplate logic user may need
|
||||
in training / testing. They will not work for everyone, but many users may find them useful.
|
||||
|
||||
The behavior of functions/classes in this file is subject to change,
|
||||
since they are meant to represent the "common default behavior" people need in their projects.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from typing import Optional
|
||||
import torch
|
||||
from fvcore.nn.precise_bn import get_bn_modules
|
||||
from omegaconf import OmegaConf
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
import detectron2.data.transforms as T
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
from detectron2.config import CfgNode, LazyConfig
|
||||
from detectron2.data import (
|
||||
MetadataCatalog,
|
||||
build_detection_test_loader,
|
||||
build_detection_train_loader,
|
||||
)
|
||||
from detectron2.evaluation import (
|
||||
DatasetEvaluator,
|
||||
inference_on_dataset,
|
||||
print_csv_format,
|
||||
verify_results,
|
||||
)
|
||||
from detectron2.modeling import build_model
|
||||
from detectron2.solver import build_lr_scheduler, build_optimizer
|
||||
from detectron2.utils import comm
|
||||
from detectron2.utils.collect_env import collect_env_info
|
||||
from detectron2.utils.env import seed_all_rng
|
||||
from detectron2.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
|
||||
from detectron2.utils.file_io import PathManager
|
||||
from detectron2.utils.logger import setup_logger
|
||||
|
||||
from detectron2.engine import hooks
|
||||
from detectron2.engine.train_loop import AMPTrainer, SimpleTrainer, TrainerBase
|
||||
|
||||
from .mycheckpointer import MyDetectionCheckpointer
|
||||
from typing import Any, Dict, List, Set
|
||||
import itertools
|
||||
from detectron2.solver.build import maybe_add_gradient_clipping
|
||||
from .dataset_mapper import DetrDatasetMapper
|
||||
from .funsd_evaluation import FUNSDEvaluator
|
||||
|
||||
__all__ = [
|
||||
"create_ddp_model",
|
||||
"default_argument_parser",
|
||||
"default_setup",
|
||||
"default_writers",
|
||||
"DefaultPredictor",
|
||||
"MyTrainer",
|
||||
]
|
||||
|
||||
|
||||
def create_ddp_model(model, *, fp16_compression=False, **kwargs):
|
||||
"""
|
||||
Create a DistributedDataParallel model if there are >1 processes.
|
||||
|
||||
Args:
|
||||
model: a torch.nn.Module
|
||||
fp16_compression: add fp16 compression hooks to the ddp object.
|
||||
See more at https://pytorch.org/docs/stable/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook
|
||||
kwargs: other arguments of :module:`torch.nn.parallel.DistributedDataParallel`.
|
||||
""" # noqa
|
||||
if comm.get_world_size() == 1:
|
||||
return model
|
||||
if "device_ids" not in kwargs:
|
||||
kwargs["device_ids"] = [comm.get_local_rank()]
|
||||
ddp = DistributedDataParallel(model, **kwargs)
|
||||
if fp16_compression:
|
||||
from torch.distributed.algorithms.ddp_comm_hooks import default as comm_hooks
|
||||
|
||||
ddp.register_comm_hook(state=None, hook=comm_hooks.fp16_compress_hook)
|
||||
return ddp
|
||||
|
||||
|
||||
def default_argument_parser(epilog=None):
|
||||
"""
|
||||
Create a parser with some common arguments used by detectron2 users.
|
||||
|
||||
Args:
|
||||
epilog (str): epilog passed to ArgumentParser describing the usage.
|
||||
|
||||
Returns:
|
||||
argparse.ArgumentParser:
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
epilog=epilog
|
||||
or f"""
|
||||
Examples:
|
||||
|
||||
Run on single machine:
|
||||
$ {sys.argv[0]} --num-gpus 8 --config-file cfg.yaml
|
||||
|
||||
Change some config options:
|
||||
$ {sys.argv[0]} --config-file cfg.yaml MODEL.WEIGHTS /path/to/weight.pth SOLVER.BASE_LR 0.001
|
||||
|
||||
Run on multiple machines:
|
||||
(machine0)$ {sys.argv[0]} --machine-rank 0 --num-machines 2 --dist-url <URL> [--other-flags]
|
||||
(machine1)$ {sys.argv[0]} --machine-rank 1 --num-machines 2 --dist-url <URL> [--other-flags]
|
||||
""",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
|
||||
parser.add_argument(
|
||||
"--resume",
|
||||
action="store_true",
|
||||
help="Whether to attempt to resume from the checkpoint directory. "
|
||||
"See documentation of `MyTrainer.resume_or_load()` for what it means.",
|
||||
)
|
||||
parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
|
||||
parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")
|
||||
parser.add_argument("--num-machines", type=int, default=1, help="total number of machines")
|
||||
parser.add_argument(
|
||||
"--machine-rank", type=int, default=0, help="the rank of this machine (unique per machine)"
|
||||
)
|
||||
|
||||
# PyTorch still may leave orphan processes in multi-gpu training.
|
||||
# Therefore we use a deterministic way to obtain port,
|
||||
# so that users are aware of orphan processes by seeing the port occupied.
|
||||
port = 2 ** 15 + 2 ** 14 + hash(os.getuid() if sys.platform != "win32" else 1) % 2 ** 14
|
||||
parser.add_argument(
|
||||
"--dist-url",
|
||||
default="tcp://127.0.0.1:{}".format(port),
|
||||
help="initialization URL for pytorch distributed backend. See "
|
||||
"https://pytorch.org/docs/stable/distributed.html for details.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"opts",
|
||||
help="""
|
||||
Modify config options at the end of the command. For Yacs configs, use
|
||||
space-separated "PATH.KEY VALUE" pairs.
|
||||
For python-based LazyConfig, use "path.key=value".
|
||||
""".strip(),
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def _try_get_key(cfg, *keys, default=None):
|
||||
"""
|
||||
Try select keys from cfg until the first key that exists. Otherwise return default.
|
||||
"""
|
||||
if isinstance(cfg, CfgNode):
|
||||
cfg = OmegaConf.create(cfg.dump())
|
||||
for k in keys:
|
||||
none = object()
|
||||
p = OmegaConf.select(cfg, k, default=none)
|
||||
if p is not none:
|
||||
return p
|
||||
return default
|
||||
|
||||
|
||||
def _highlight(code, filename):
|
||||
try:
|
||||
import pygments
|
||||
except ImportError:
|
||||
return code
|
||||
|
||||
from pygments.lexers import Python3Lexer, YamlLexer
|
||||
from pygments.formatters import Terminal256Formatter
|
||||
|
||||
lexer = Python3Lexer() if filename.endswith(".py") else YamlLexer()
|
||||
code = pygments.highlight(code, lexer, Terminal256Formatter(style="monokai"))
|
||||
return code
|
||||
|
||||
|
||||
def default_setup(cfg, args):
|
||||
"""
|
||||
Perform some basic common setups at the beginning of a job, including:
|
||||
|
||||
1. Set up the detectron2 logger
|
||||
2. Log basic information about environment, cmdline arguments, and config
|
||||
3. Backup the config to the output directory
|
||||
|
||||
Args:
|
||||
cfg (CfgNode or omegaconf.DictConfig): the full config to be used
|
||||
args (argparse.NameSpace): the command line arguments to be logged
|
||||
"""
|
||||
output_dir = _try_get_key(cfg, "OUTPUT_DIR", "output_dir", "train.output_dir")
|
||||
if comm.is_main_process() and output_dir:
|
||||
PathManager.mkdirs(output_dir)
|
||||
|
||||
rank = comm.get_rank()
|
||||
setup_logger(output_dir, distributed_rank=rank, name="fvcore")
|
||||
logger = setup_logger(output_dir, distributed_rank=rank)
|
||||
|
||||
logger.info("Rank of current process: {}. World size: {}".format(rank, comm.get_world_size()))
|
||||
logger.info("Environment info:\n" + collect_env_info())
|
||||
|
||||
logger.info("Command line arguments: " + str(args))
|
||||
if hasattr(args, "config_file") and args.config_file != "":
|
||||
logger.info(
|
||||
"Contents of args.config_file={}:\n{}".format(
|
||||
args.config_file,
|
||||
_highlight(PathManager.open(args.config_file, "r").read(), args.config_file),
|
||||
)
|
||||
)
|
||||
|
||||
if comm.is_main_process() and output_dir:
|
||||
# Note: some of our scripts may expect the existence of
|
||||
# config.yaml in output directory
|
||||
path = os.path.join(output_dir, "config.yaml")
|
||||
if isinstance(cfg, CfgNode):
|
||||
logger.info("Running with full config:\n{}".format(_highlight(cfg.dump(), ".yaml")))
|
||||
with PathManager.open(path, "w") as f:
|
||||
f.write(cfg.dump())
|
||||
else:
|
||||
LazyConfig.save(cfg, path)
|
||||
logger.info("Full config saved to {}".format(path))
|
||||
|
||||
# make sure each worker has a different, yet deterministic seed if specified
|
||||
seed = _try_get_key(cfg, "SEED", "train.seed", default=-1)
|
||||
seed_all_rng(None if seed < 0 else seed + rank)
|
||||
|
||||
# cudnn benchmark has large overhead. It shouldn't be used considering the small size of
|
||||
# typical validation set.
|
||||
if not (hasattr(args, "eval_only") and args.eval_only):
|
||||
torch.backends.cudnn.benchmark = _try_get_key(
|
||||
cfg, "CUDNN_BENCHMARK", "train.cudnn_benchmark", default=False
|
||||
)
|
||||
|
||||
|
||||
def default_writers(output_dir: str, max_iter: Optional[int] = None):
|
||||
"""
|
||||
Build a list of :class:`EventWriter` to be used.
|
||||
It now consists of a :class:`CommonMetricPrinter`,
|
||||
:class:`TensorboardXWriter` and :class:`JSONWriter`.
|
||||
|
||||
Args:
|
||||
output_dir: directory to store JSON metrics and tensorboard events
|
||||
max_iter: the total number of iterations
|
||||
|
||||
Returns:
|
||||
list[EventWriter]: a list of :class:`EventWriter` objects.
|
||||
"""
|
||||
PathManager.mkdirs(output_dir)
|
||||
return [
|
||||
# It may not always print what you want to see, since it prints "common" metrics only.
|
||||
CommonMetricPrinter(max_iter),
|
||||
JSONWriter(os.path.join(output_dir, "metrics.json")),
|
||||
TensorboardXWriter(output_dir),
|
||||
]
|
||||
|
||||
|
||||
class DefaultPredictor:
|
||||
"""
|
||||
Create a simple end-to-end predictor with the given config that runs on
|
||||
single device for a single input image.
|
||||
|
||||
Compared to using the model directly, this class does the following additions:
|
||||
|
||||
1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
|
||||
2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
|
||||
3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
|
||||
4. Take one input image and produce a single output, instead of a batch.
|
||||
|
||||
This is meant for simple demo purposes, so it does the above steps automatically.
|
||||
This is not meant for benchmarks or running complicated inference logic.
|
||||
If you'd like to do anything more complicated, please refer to its source code as
|
||||
examples to build and use the model manually.
|
||||
|
||||
Attributes:
|
||||
metadata (Metadata): the metadata of the underlying dataset, obtained from
|
||||
cfg.DATASETS.TEST.
|
||||
|
||||
Examples:
|
||||
::
|
||||
pred = DefaultPredictor(cfg)
|
||||
inputs = cv2.imread("input.jpg")
|
||||
outputs = pred(inputs)
|
||||
"""
|
||||
|
||||
def __init__(self, cfg):
|
||||
self.cfg = cfg.clone() # cfg can be modified by model
|
||||
self.model = build_model(self.cfg)
|
||||
self.model.eval()
|
||||
if len(cfg.DATASETS.TEST):
|
||||
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
|
||||
|
||||
checkpointer = DetectionCheckpointer(self.model)
|
||||
checkpointer.load(cfg.MODEL.WEIGHTS)
|
||||
|
||||
self.aug = T.ResizeShortestEdge(
|
||||
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
|
||||
)
|
||||
|
||||
self.input_format = cfg.INPUT.FORMAT
|
||||
assert self.input_format in ["RGB", "BGR"], self.input_format
|
||||
|
||||
def __call__(self, original_image):
|
||||
"""
|
||||
Args:
|
||||
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
|
||||
|
||||
Returns:
|
||||
predictions (dict):
|
||||
the output of the model for one image only.
|
||||
See :doc:`/tutorials/models` for details about the format.
|
||||
"""
|
||||
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
|
||||
# Apply pre-processing to image.
|
||||
if self.input_format == "RGB":
|
||||
# whether the model expects BGR inputs or RGB
|
||||
original_image = original_image[:, :, ::-1]
|
||||
height, width = original_image.shape[:2]
|
||||
image = self.aug.get_transform(original_image).apply_image(original_image)
|
||||
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
|
||||
|
||||
inputs = {"image": image, "height": height, "width": width}
|
||||
predictions = self.model([inputs])[0]
|
||||
return predictions
|
||||
|
||||
|
||||
class MyTrainer(TrainerBase):
|
||||
"""
|
||||
A trainer with default training logic. It does the following:
|
||||
|
||||
1. Create a :class:`SimpleTrainer` using model, optimizer, dataloader
|
||||
defined by the given config. Create a LR scheduler defined by the config.
|
||||
2. Load the last checkpoint or `cfg.MODEL.WEIGHTS`, if exists, when
|
||||
`resume_or_load` is called.
|
||||
3. Register a few common hooks defined by the config.
|
||||
|
||||
It is created to simplify the **standard model training workflow** and reduce code boilerplate
|
||||
for users who only need the standard training workflow, with standard features.
|
||||
It means this class makes *many assumptions* about your training logic that
|
||||
may easily become invalid in a new research. In fact, any assumptions beyond those made in the
|
||||
:class:`SimpleTrainer` are too much for research.
|
||||
|
||||
The code of this class has been annotated about restrictive assumptions it makes.
|
||||
When they do not work for you, you're encouraged to:
|
||||
|
||||
1. Overwrite methods of this class, OR:
|
||||
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
|
||||
nothing else. You can then add your own hooks if needed. OR:
|
||||
3. Write your own training loop similar to `tools/plain_train_net.py`.
|
||||
|
||||
See the :doc:`/tutorials/training` tutorials for more details.
|
||||
|
||||
Note that the behavior of this class, like other functions/classes in
|
||||
this file, is not stable, since it is meant to represent the "common default behavior".
|
||||
It is only guaranteed to work well with the standard models and training workflow in detectron2.
|
||||
To obtain more stable behavior, write your own training logic with other public APIs.
|
||||
|
||||
Examples:
|
||||
::
|
||||
trainer = MyTrainer(cfg)
|
||||
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
|
||||
trainer.train()
|
||||
|
||||
Attributes:
|
||||
scheduler:
|
||||
checkpointer (DetectionCheckpointer):
|
||||
cfg (CfgNode):
|
||||
"""
|
||||
|
||||
def __init__(self, cfg):
|
||||
"""
|
||||
Args:
|
||||
cfg (CfgNode):
|
||||
"""
|
||||
super().__init__()
|
||||
logger = logging.getLogger("detectron2")
|
||||
if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2
|
||||
setup_logger()
|
||||
cfg = MyTrainer.auto_scale_workers(cfg, comm.get_world_size())
|
||||
|
||||
self.cfg = cfg
|
||||
|
||||
# Assume these objects must be constructed in this order.
|
||||
model = self.build_model(cfg)
|
||||
optimizer = self.build_optimizer(cfg, model)
|
||||
data_loader = self.build_train_loader(cfg)
|
||||
|
||||
model = create_ddp_model(model, broadcast_buffers=False)
|
||||
self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
|
||||
model, data_loader, optimizer
|
||||
)
|
||||
|
||||
self.scheduler = self.build_lr_scheduler(cfg, optimizer)
|
||||
self.checkpointer = MyDetectionCheckpointer(
|
||||
# Assume you want to save checkpoints together with logs/statistics
|
||||
model,
|
||||
cfg.OUTPUT_DIR,
|
||||
trainer=weakref.proxy(self),
|
||||
)
|
||||
self.start_iter = 0
|
||||
self.max_iter = cfg.SOLVER.MAX_ITER
|
||||
self.cfg = cfg
|
||||
|
||||
self.register_hooks(self.build_hooks())
|
||||
|
||||
def resume_or_load(self, resume=True):
|
||||
"""
|
||||
If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
|
||||
a `last_checkpoint` file), resume from the file. Resuming means loading all
|
||||
available states (eg. optimizer and scheduler) and update iteration counter
|
||||
from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
|
||||
|
||||
Otherwise, this is considered as an independent training. The method will load model
|
||||
weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
|
||||
from iteration 0.
|
||||
|
||||
Args:
|
||||
resume (bool): whether to do resume or not
|
||||
"""
|
||||
self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
|
||||
if resume and self.checkpointer.has_checkpoint():
|
||||
# The checkpoint stores the training iteration that just finished, thus we start
|
||||
# at the next iteration
|
||||
self.start_iter = self.iter + 1
|
||||
|
||||
def build_hooks(self):
|
||||
"""
|
||||
Build a list of default hooks, including timing, evaluation,
|
||||
checkpointing, lr scheduling, precise BN, writing events.
|
||||
|
||||
Returns:
|
||||
list[HookBase]:
|
||||
"""
|
||||
cfg = self.cfg.clone()
|
||||
cfg.defrost()
|
||||
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
|
||||
|
||||
ret = [
|
||||
hooks.IterationTimer(),
|
||||
hooks.LRScheduler(),
|
||||
hooks.PreciseBN(
|
||||
# Run at the same freq as (but before) evaluation.
|
||||
cfg.TEST.EVAL_PERIOD,
|
||||
self.model,
|
||||
# Build a new data loader to not affect training
|
||||
self.build_train_loader(cfg),
|
||||
cfg.TEST.PRECISE_BN.NUM_ITER,
|
||||
)
|
||||
if cfg.TEST.PRECISE_BN.ENABLED and get_bn_modules(self.model)
|
||||
else None,
|
||||
]
|
||||
|
||||
# Do PreciseBN before checkpointer, because it updates the model and need to
|
||||
# be saved by checkpointer.
|
||||
# This is not always the best: if checkpointing has a different frequency,
|
||||
# some checkpoints may have more precise statistics than others.
|
||||
if comm.is_main_process():
|
||||
ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
|
||||
|
||||
def test_and_save_results():
|
||||
self._last_eval_results = self.test(self.cfg, self.model)
|
||||
return self._last_eval_results
|
||||
|
||||
# Do evaluation after checkpointer, because then if it fails,
|
||||
# we can use the saved checkpoint to debug.
|
||||
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
|
||||
|
||||
if comm.is_main_process():
|
||||
# Here the default print/log frequency of each writer is used.
|
||||
# run writers in the end, so that evaluation metrics are written
|
||||
ret.append(hooks.PeriodicWriter(self.build_writers(), period=20))
|
||||
return ret
|
||||
|
||||
def build_writers(self):
|
||||
"""
|
||||
Build a list of writers to be used using :func:`default_writers()`.
|
||||
If you'd like a different list of writers, you can overwrite it in
|
||||
your trainer.
|
||||
|
||||
Returns:
|
||||
list[EventWriter]: a list of :class:`EventWriter` objects.
|
||||
"""
|
||||
return default_writers(self.cfg.OUTPUT_DIR, self.max_iter)
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
Run training.
|
||||
|
||||
Returns:
|
||||
OrderedDict of results, if evaluation is enabled. Otherwise None.
|
||||
"""
|
||||
super().train(self.start_iter, self.max_iter)
|
||||
if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
|
||||
assert hasattr(
|
||||
self, "_last_eval_results"
|
||||
), "No evaluation results obtained during training!"
|
||||
verify_results(self.cfg, self._last_eval_results)
|
||||
return self._last_eval_results
|
||||
|
||||
def run_step(self):
|
||||
self._trainer.iter = self.iter
|
||||
self._trainer.run_step()
|
||||
|
||||
@classmethod
|
||||
def build_model(cls, cfg):
|
||||
"""
|
||||
Returns:
|
||||
torch.nn.Module:
|
||||
|
||||
It now calls :func:`detectron2.modeling.build_model`.
|
||||
Overwrite it if you'd like a different model.
|
||||
"""
|
||||
model = build_model(cfg)
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info("Model:\n{}".format(model))
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def build_optimizer(cls, cfg, model):
|
||||
params: List[Dict[str, Any]] = []
|
||||
memo: Set[torch.nn.parameter.Parameter] = set()
|
||||
for key, value in model.named_parameters(recurse=True):
|
||||
if not value.requires_grad:
|
||||
continue
|
||||
# Avoid duplicating parameters
|
||||
if value in memo:
|
||||
continue
|
||||
memo.add(value)
|
||||
lr = cfg.SOLVER.BASE_LR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY
|
||||
if "backbone" in key:
|
||||
lr = lr * cfg.SOLVER.BACKBONE_MULTIPLIER
|
||||
params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
|
||||
|
||||
def maybe_add_full_model_gradient_clipping(optim): # optim: the optimizer class
|
||||
# detectron2 doesn't have full model gradient clipping now
|
||||
clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
|
||||
enable = (
|
||||
cfg.SOLVER.CLIP_GRADIENTS.ENABLED
|
||||
and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
|
||||
and clip_norm_val > 0.0
|
||||
)
|
||||
|
||||
class FullModelGradientClippingOptimizer(optim):
|
||||
def step(self, closure=None):
|
||||
all_params = itertools.chain(*[x["params"] for x in self.param_groups])
|
||||
torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
|
||||
super().step(closure=closure)
|
||||
|
||||
return FullModelGradientClippingOptimizer if enable else optim
|
||||
|
||||
optimizer_type = cfg.SOLVER.OPTIMIZER
|
||||
if optimizer_type == "SGD":
|
||||
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
|
||||
params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
|
||||
)
|
||||
elif optimizer_type == "ADAMW":
|
||||
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
|
||||
params, cfg.SOLVER.BASE_LR
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"no optimizer type {optimizer_type}")
|
||||
if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
|
||||
optimizer = maybe_add_gradient_clipping(cfg, optimizer)
|
||||
return optimizer
|
||||
|
||||
@classmethod
|
||||
def build_lr_scheduler(cls, cfg, optimizer):
|
||||
"""
|
||||
It now calls :func:`detectron2.solver.build_lr_scheduler`.
|
||||
Overwrite it if you'd like a different scheduler.
|
||||
"""
|
||||
return build_lr_scheduler(cfg, optimizer)
|
||||
|
||||
@classmethod
|
||||
def build_train_loader(cls, cfg):
|
||||
if cfg.AUG.DETR:
|
||||
mapper = DetrDatasetMapper(cfg, is_train=True)
|
||||
else:
|
||||
mapper = None
|
||||
return build_detection_train_loader(cfg, mapper=mapper)
|
||||
|
||||
@classmethod
|
||||
def build_test_loader(cls, cfg, dataset_name):
|
||||
"""
|
||||
Returns:
|
||||
iterable
|
||||
|
||||
It now calls :func:`detectron2.data.build_detection_test_loader`.
|
||||
Overwrite it if you'd like a different data loader.
|
||||
"""
|
||||
return build_detection_test_loader(cfg, dataset_name)
|
||||
|
||||
@classmethod
|
||||
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
|
||||
if output_folder is None:
|
||||
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
|
||||
return FUNSDEvaluator(dataset_name, output_dir=output_folder)
|
||||
|
||||
@classmethod
|
||||
def test(cls, cfg, model, evaluators=None):
|
||||
"""
|
||||
Evaluate the given model. The given model is expected to already contain
|
||||
weights to evaluate.
|
||||
|
||||
Args:
|
||||
cfg (CfgNode):
|
||||
model (nn.Module):
|
||||
evaluators (list[DatasetEvaluator] or None): if None, will call
|
||||
:meth:`build_evaluator`. Otherwise, must have the same length as
|
||||
``cfg.DATASETS.TEST``.
|
||||
|
||||
Returns:
|
||||
dict: a dict of result metrics
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
if isinstance(evaluators, DatasetEvaluator):
|
||||
evaluators = [evaluators]
|
||||
if evaluators is not None:
|
||||
assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
|
||||
len(cfg.DATASETS.TEST), len(evaluators)
|
||||
)
|
||||
|
||||
results = OrderedDict()
|
||||
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
|
||||
data_loader = cls.build_test_loader(cfg, dataset_name)
|
||||
# When evaluators are passed in as arguments,
|
||||
# implicitly assume that evaluators can be created before data_loader.
|
||||
if evaluators is not None:
|
||||
evaluator = evaluators[idx]
|
||||
else:
|
||||
try:
|
||||
evaluator = cls.build_evaluator(cfg, dataset_name)
|
||||
except NotImplementedError:
|
||||
logger.warn(
|
||||
"No evaluator found. Use `MyTrainer.test(evaluators=)`, "
|
||||
"or implement its `build_evaluator` method."
|
||||
)
|
||||
results[dataset_name] = {}
|
||||
continue
|
||||
results_i = inference_on_dataset(model, data_loader, evaluator)
|
||||
results[dataset_name] = results_i
|
||||
if comm.is_main_process():
|
||||
assert isinstance(
|
||||
results_i, dict
|
||||
), "Evaluator must return a dict on the main process. Got {} instead.".format(
|
||||
results_i
|
||||
)
|
||||
logger.info("Evaluation results for {} in csv format:".format(dataset_name))
|
||||
print_csv_format(results_i)
|
||||
|
||||
if len(results) == 1:
|
||||
results = list(results.values())[0]
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def auto_scale_workers(cfg, num_workers: int):
|
||||
"""
|
||||
When the config is defined for certain number of workers (according to
|
||||
``cfg.SOLVER.REFERENCE_WORLD_SIZE``) that's different from the number of
|
||||
workers currently in use, returns a new cfg where the total batch size
|
||||
is scaled so that the per-GPU batch size stays the same as the
|
||||
original ``IMS_PER_BATCH // REFERENCE_WORLD_SIZE``.
|
||||
|
||||
Other config options are also scaled accordingly:
|
||||
* training steps and warmup steps are scaled inverse proportionally.
|
||||
* learning rate are scaled proportionally, following :paper:`ImageNet in 1h`.
|
||||
|
||||
For example, with the original config like the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
IMS_PER_BATCH: 16
|
||||
BASE_LR: 0.1
|
||||
REFERENCE_WORLD_SIZE: 8
|
||||
MAX_ITER: 5000
|
||||
STEPS: (4000,)
|
||||
CHECKPOINT_PERIOD: 1000
|
||||
|
||||
When this config is used on 16 GPUs instead of the reference number 8,
|
||||
calling this method will return a new config with:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
IMS_PER_BATCH: 32
|
||||
BASE_LR: 0.2
|
||||
REFERENCE_WORLD_SIZE: 16
|
||||
MAX_ITER: 2500
|
||||
STEPS: (2000,)
|
||||
CHECKPOINT_PERIOD: 500
|
||||
|
||||
Note that both the original config and this new config can be trained on 16 GPUs.
|
||||
It's up to user whether to enable this feature (by setting ``REFERENCE_WORLD_SIZE``).
|
||||
|
||||
Returns:
|
||||
CfgNode: a new config. Same as original if ``cfg.SOLVER.REFERENCE_WORLD_SIZE==0``.
|
||||
"""
|
||||
old_world_size = cfg.SOLVER.REFERENCE_WORLD_SIZE
|
||||
if old_world_size == 0 or old_world_size == num_workers:
|
||||
return cfg
|
||||
cfg = cfg.clone()
|
||||
frozen = cfg.is_frozen()
|
||||
cfg.defrost()
|
||||
|
||||
assert (
|
||||
cfg.SOLVER.IMS_PER_BATCH % old_world_size == 0
|
||||
), "Invalid REFERENCE_WORLD_SIZE in config!"
|
||||
scale = num_workers / old_world_size
|
||||
bs = cfg.SOLVER.IMS_PER_BATCH = int(round(cfg.SOLVER.IMS_PER_BATCH * scale))
|
||||
lr = cfg.SOLVER.BASE_LR = cfg.SOLVER.BASE_LR * scale
|
||||
max_iter = cfg.SOLVER.MAX_ITER = int(round(cfg.SOLVER.MAX_ITER / scale))
|
||||
warmup_iter = cfg.SOLVER.WARMUP_ITERS = int(round(cfg.SOLVER.WARMUP_ITERS / scale))
|
||||
cfg.SOLVER.STEPS = tuple(int(round(s / scale)) for s in cfg.SOLVER.STEPS)
|
||||
cfg.TEST.EVAL_PERIOD = int(round(cfg.TEST.EVAL_PERIOD / scale))
|
||||
cfg.SOLVER.CHECKPOINT_PERIOD = int(round(cfg.SOLVER.CHECKPOINT_PERIOD / scale))
|
||||
cfg.SOLVER.REFERENCE_WORLD_SIZE = num_workers # maintain invariant
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(
|
||||
f"Auto-scaling the config to batch_size={bs}, learning_rate={lr}, "
|
||||
f"max_iter={max_iter}, warmup={warmup_iter}."
|
||||
)
|
||||
|
||||
if frozen:
|
||||
cfg.freeze()
|
||||
return cfg
|
||||
|
||||
|
||||
# Access basic attributes from the underlying trainer
|
||||
for _attr in ["model", "data_loader", "optimizer"]:
|
||||
setattr(
|
||||
MyTrainer,
|
||||
_attr,
|
||||
property(
|
||||
# getter
|
||||
lambda self, x=_attr: getattr(self._trainer, x),
|
||||
# setter
|
||||
lambda self, value, x=_attr: setattr(self._trainer, x, value),
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,84 @@
|
||||
#!/usr/bin/env python
|
||||
# --------------------------------------------------------------------------------
|
||||
# MPViT: Multi-Path Vision Transformer for Dense Prediction
|
||||
# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
|
||||
# All Rights Reserved.
|
||||
# Written by Youngwan Lee
|
||||
# --------------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Detection Training Script for MPViT.
|
||||
"""
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
from detectron2.config import get_cfg
|
||||
from detectron2.engine import default_argument_parser, default_setup, launch
|
||||
from detectron2.data.datasets import register_coco_instances
|
||||
|
||||
from ditod import MyTrainer, add_vit_config
|
||||
|
||||
def setup(args):
|
||||
"""
|
||||
Create configs and perform basic setups.
|
||||
"""
|
||||
cfg = get_cfg()
|
||||
# add_coat_config(cfg)
|
||||
add_vit_config(cfg)
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
default_setup(cfg, args)
|
||||
return cfg
|
||||
|
||||
|
||||
def main(args):
|
||||
register_coco_instances(
|
||||
"funsd_train",
|
||||
{},
|
||||
"data/instances_training.json",
|
||||
"data/imgs"
|
||||
)
|
||||
|
||||
register_coco_instances(
|
||||
"funsd_test",
|
||||
{},
|
||||
"data/instances_test.json",
|
||||
"data/imgs"
|
||||
)
|
||||
|
||||
cfg = setup(args)
|
||||
|
||||
if args.eval_only:
|
||||
model = MyTrainer.build_model(cfg)
|
||||
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
|
||||
cfg.MODEL.WEIGHTS, resume=args.resume
|
||||
)
|
||||
res = MyTrainer.test(cfg, model)
|
||||
return res
|
||||
|
||||
trainer = MyTrainer(cfg)
|
||||
trainer.resume_or_load(resume=args.resume)
|
||||
return trainer.train()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = default_argument_parser()
|
||||
parser.add_argument("--debug", action="store_true", help="enable debug mode")
|
||||
args = parser.parse_args()
|
||||
print("Command Line Args:", args)
|
||||
|
||||
if args.debug:
|
||||
import debugpy
|
||||
|
||||
print("Enabling attach starts.")
|
||||
debugpy.listen(address=('0.0.0.0', 9310))
|
||||
debugpy.wait_for_client()
|
||||
print("Enabling attach ends.")
|
||||
|
||||
launch(
|
||||
main,
|
||||
args.num_gpus,
|
||||
num_machines=args.num_machines,
|
||||
machine_rank=args.machine_rank,
|
||||
dist_url=args.dist_url,
|
||||
args=(args,),
|
||||
)
|
||||
Reference in New Issue
Block a user