Initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+deeplabv3_mobilenet_v3_large filter=lfs diff=lfs merge=lfs -text

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "vision"]
+	path = vision
+	url = https://github.com/pytorch/vision
+	shallow = true

common.py

@@ -0,0 +1,84 @@
+import os
+
+import torch
+import torchvision
+from fvcore.nn import FlopCountAnalysis
+from torch import nn
+from transforms import Compose
+
+sys.path.append("vision/references/segmentation")
+from coco_utils import ConvertCocoPolysToMask
+from coco_utils import FilterAndRemapCocoCategories
+from coco_utils import _coco_remove_images_without_annotations
+from utils import ConfusionMatrix
+
+
+class NanSafeConfusionMatrix(ConfusionMatrix):
+    """Confusion matrix with replacement nans to zeros."""
+
+    def __init__(self, num_classes):
+        super().__init__(num_classes=num_classes)
+
+    def compute(self):
+        """Compute metrics based on confusion matrix."""
+        confusion_matrix = self.mat.float()
+        acc_global = torch.nan_to_num(torch.diag(confusion_matrix).sum() / confusion_matrix.sum())
+        acc = torch.nan_to_num(torch.diag(confusion_matrix) / confusion_matrix.sum(1))
+        intersection_over_unions = torch.nan_to_num(
+            torch.diag(confusion_matrix)
+            / (confusion_matrix.sum(1) + confusion_matrix.sum(0) - torch.diag(confusion_matrix))
+        )
+        return acc_global, acc, intersection_over_unions
+
+
+def flops_calculation_function(model: nn.Module, input_sample: torch.Tensor) -> float:
+    """Calculate number of flops in millions."""
+    counter = FlopCountAnalysis(
+        model=model.eval(),
+        inputs=input_sample,
+    )
+    counter.unsupported_ops_warnings(False)
+    counter.uncalled_modules_warnings(False)
+
+    flops = counter.total() / input_sample.shape[0]
+
+    return flops / 1e6
+
+
+def get_coco(root, image_set, transforms, use_v2=False, use_orig=False):
+    """Get COCO dataset with VOC or COCO classes."""
+    paths = {
+        "train": ("train2017", os.path.join("annotations", "instances_train2017.json")),
+        "val": ("val2017", os.path.join("annotations", "instances_val2017.json")),
+        # "train": ("val2017", os.path.join("annotations", "instances_val2017.json"))
+    }
+    if use_orig:
+        classes_list = list(range(81))
+    else:
+        classes_list = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4, 1, 64, 20, 63, 7, 72]
+
+    img_folder, ann_file = paths[image_set]
+    img_folder = os.path.join(root, img_folder)
+    ann_file = os.path.join(root, ann_file)
+
+    # The 2 "Compose" below achieve the same thing: converting coco detection
+    # samples into segmentation-compatible samples. They just do it with
+    # slightly different implementations. We could refactor and unify, but
+    # keeping them separate helps keeping the v2 version clean
+    if use_v2:
+        import v2_extras  # pylint: disable=import-outside-toplevel
+        from torchvision.datasets import wrap_dataset_for_transforms_v2  # pylint: disable=import-outside-toplevel
+
+        transforms = Compose([v2_extras.CocoDetectionToVOCSegmentation(), transforms])
+        dataset = torchvision.datasets.CocoDetection(img_folder, ann_file, transforms=transforms)
+        dataset = wrap_dataset_for_transforms_v2(dataset, target_keys={"masks", "labels"})
+    else:
+        transforms = Compose(
+            [FilterAndRemapCocoCategories(classes_list, remap=True), ConvertCocoPolysToMask(), transforms]
+        )
+        dataset = torchvision.datasets.CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    if image_set == "train":
+        dataset = _coco_remove_images_without_annotations(dataset, classes_list)
+
+    return dataset

deeplabv3_mobilenet_v3_large/deeplabv3_mobilenet_v3_large.pth

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

deeplabv3_mobilenet_v3_large/deeplabv3_mobilenet_v3_large_x2.pth

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

deeplabv3_mobilenet_v3_large/deeplabv3_mobilenet_v3_large_x4.pth

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

requirements.txt

@@ -0,0 +1,2 @@
+torch==2.3.1
+torchvision==0.18.1

test.py

@@ -0,0 +1,165 @@
+import sys
+from functools import partial
+
+from typing import Callable
+from typing import Dict
+from typing import Tuple
+from typing import Union
+from argparse import Namespace
+
+sys.path.append("vision/references/segmentation")
+
+import presets
+import torch
+import torch.utils.data
+import torchvision
+import utils
+from torch import nn
+from common import flops_calculation_function
+from common import NanSafeConfusionMatrix as ConfusionMatrix
+from common import get_coco
+
+
+def get_dataset(args: Namespace, is_train: bool, transform: Callable = None) -> Tuple[torch.utils.data.Dataset, int]:
+    def sbd(*args, **kwargs):
+        kwargs.pop("use_v2")
+        return torchvision.datasets.SBDataset(*args, mode="segmentation", **kwargs)
+
+    def voc(*args, **kwargs):
+        kwargs.pop("use_v2")
+        return torchvision.datasets.VOCSegmentation(*args, **kwargs)
+
+    paths = {
+        "voc": (args.data_path, voc, 21),
+        "voc_aug": (args.data_path, sbd, 21),
+        "coco": (args.data_path, get_coco, 21),
+        "coco_orig": (args.data_path, partial(get_coco, use_orig=True), 81)
+    }
+    p, ds_fn, num_classes = paths["coco_orig"]
+
+    if transform is None:
+        transform = get_transform(is_train, args)
+    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():
+        losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
+
+    if len(losses) == 1:
+        return losses["out"]
+
+    return losses["out"] + 0.5 * losses["aux"]
+
+
+def evaluate(
+        model: torch.nn.Module,
+        data_loader: torch.utils.data.DataLoader,
+        device: Union[str, torch.device],
+        num_classes: int,
+        criterion: Callable,
+) -> Tuple[ConfusionMatrix, float]:
+    model.eval()
+    confmat = ConfusionMatrix(num_classes)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = "Test:"
+    num_processed_samples = 0
+    with torch.inference_mode():
+        for batch_n, (image, target) in enumerate(metric_logger.log_every(data_loader, 100, header)):
+            image, target = image.to(device), target.to(device)
+            output = model(image)
+            loss = criterion(output, target)
+            output = output["out"]
+
+            confmat.update(target.flatten(), output.argmax(1).flatten())
+            # FIXME need to take into account that the datasets
+            # could have been padded in distributed setup
+            num_processed_samples += image.shape[0]
+
+            metric_logger.update(loss=loss.item())
+
+        confmat.reduce_from_all_processes()
+
+    return confmat, metric_logger.loss.global_avg
+
+
+def main(args):
+    if args.backend.lower() != "pil" and not args.use_v2:
+        # TODO: Support tensor backend in V1?
+        raise ValueError("Use --use-v2 if you want to use the tv_tensor or tensor backend.")
+    if args.use_v2:
+        raise ValueError("v2 is only supported for coco dataset for now.")
+
+    print(args)
+
+    device = torch.device(args.device)
+
+    if args.use_deterministic_algorithms:
+        torch.backends.cudnn.benchmark = False
+        torch.use_deterministic_algorithms(True)
+    else:
+        torch.backends.cudnn.benchmark = True
+
+    dataset_test, num_classes = get_dataset(args, is_train=False)
+
+    test_sampler = torch.utils.data.SequentialSampler(dataset_test)
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
+
+    checkpoint = torch.load(args.model_path)
+    model = checkpoint["model"]
+    model.to(device)
+    model_flops = flops_calculation_function(model=model, input_sample=next(iter(data_loader_test))[0].to(device))
+    print(f"Model Flops: {model_flops}M")
+
+    # We disable the cudnn benchmarking because it can noticeably affect the accuracy
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    confmat, loss = evaluate(
+        model=model,
+        data_loader=data_loader_test,
+        device=device,
+        num_classes=num_classes,
+        criterion=criterion,
+    )
+    print(confmat)
+    return
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(description="PyTorch Segmentation Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="/datasets01/COCO/022719/", type=str, help="dataset path")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument(
+        "-b", "--batch-size", default=8, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
+    )
+    parser.add_argument("--epochs", default=30, type=int, metavar="N", help="number of total epochs to run")
+
+    parser.add_argument(
+        "-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)"
+    )
+    parser.add_argument(
+        "--use-deterministic-algorithms", action="store_true", help="Forces the use of deterministic algorithms only."
+    )
+    # distributed training parameters
+
+    parser.add_argument("--backend", default="PIL", type=str.lower, help="PIL or tensor - case insensitive")
+    parser.add_argument("--use-v2", action="store_true", help="Use V2 transforms")
+    parser.add_argument("--model-path", default=None, help="Path to model checkpoint.")
+    return parser
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args)

vision

@@ -0,0 +1 @@
+Subproject commit 126fc22ce33e6c2426edcf9ed540810c178fe9ce