Add README and mac measuring

README.md

@@ -1,3 +1,43 @@
 ---
 license: apache-2.0
+datasets:
+- coco
+pipeline_tag: image-segmentation
+tags:
+- computer-vision
+- image-segmentation
+- ENOT-AutoDL
 ---
+
+# ENOT-AutoDL pruning benchmark on MS-COCO
+
+This repository contains models accelerated with [ENOT-AutoDL](https://pypi.org/project/enot-autodl/) framework.
+Models from [Torchvision](https://pytorch.org/vision/stable/models.html) are used as a baseline.
+Evaluation code is also based on Torchvision references.
+
+## DeeplabV3_MobileNetV3_Large
+
+| Model                                       | Latency (MMACs) | mean IoU (%) |
+|---------------------------------------------|:---------------:|:------------:|
+| **DeeplabV3_MobileNetV3_Large Torchvision** |     8872.87     |     47.0     |
+| **DeeplabV3_MobileNetV3_Large ENOT (x2)**   | 4436.41 (x2.0)  | 47.6 (+0.6)  |
+| **DeeplabV3_MobileNetV3_Large ENOT (x4)**   | 2217.53 (x4.0)  | 46.4 (-0.6)  |
+
+# Validation
+
+To validate results, follow this steps:
+
+1. Install all required packages:
+   ```bash
+   pip install -r requrements.txt
+   ```
+1. Calculate model latency:
+   ```bash
+   python measure_mac.py --model-path path/to/model.pth
+   ```
+1. Measure mean IoU of PyTorch (.pth) model:
+   ```bash
+   python test.py --data-path path/to/imagenet --model-path path/to/model.pth
+   ```
+
+If you want to book a demo, please contact us: [email protected] .

common.py

@@ -1,12 +1,13 @@
 import os
+import sys
 
 import torch
 import torchvision
 from fvcore.nn import FlopCountAnalysis
 from torch import nn
-from transforms import Compose
 
 sys.path.append("vision/references/segmentation")
+from transforms import Compose
 from coco_utils import ConvertCocoPolysToMask
 from coco_utils import FilterAndRemapCocoCategories
 from coco_utils import _coco_remove_images_without_annotations

measure_mac.py

@@ -0,0 +1,20 @@
+import argparse
+import torch
+
+from common import flops_calculation_function
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        help="Path to models checkpoint (.pth file).",
+    )
+    args = parser.parse_args()
+
+    checkpoint = torch.load(args.model_path, map_location="cpu")
+    model = checkpoint["model"]
+
+    flops = flops_calculation_function(model, torch.ones(1, 3, 480, 480))
+
+    print(f"MMACs = {flops}")

test.py

@@ -38,16 +38,12 @@ def get_dataset(args: Namespace, is_train: bool, transform: Callable = None) ->
     p, ds_fn, num_classes = paths["coco_orig"]
 
     if transform is None:
-        transform = get_transform(is_train, args)
+        transform = presets.SegmentationPresetEval(base_size=520, backend=args.backend, use_v2=args.use_v2)
     image_set = "train" if is_train else "val"
     ds = ds_fn(p, image_set=image_set, transforms=transform, use_v2=args.use_v2)
     return ds, num_classes
 
 
-def get_transform(is_train: bool, args: Namespace) -> Callable:
-    return presets.SegmentationPresetEval(base_size=520, backend=args.backend, use_v2=args.use_v2)
-
-
 def criterion(inputs: Dict[str, torch.Tensor], target: Dict[str, torch.Tensor]) -> torch.Tensor:
     losses = {}
     for name, x in inputs.items():