From 71b3998ea4c3eac975aa7c78fb1475e6c385480c Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Mon, 3 Oct 2022 09:40:49 +0200 Subject: [PATCH 1/2] add vision_collate --- test/test_prototype_features.py | 207 +++++++++++++++------ torchvision/prototype/features/__init__.py | 2 + torchvision/prototype/features/_collate.py | 36 ++++ 3 files changed, 183 insertions(+), 62 deletions(-) create mode 100644 torchvision/prototype/features/_collate.py diff --git a/test/test_prototype_features.py b/test/test_prototype_features.py index 2701dd66be0..0a7cb7267b3 100644 --- a/test/test_prototype_features.py +++ b/test/test_prototype_features.py @@ -1,5 +1,15 @@ +import functools + import pytest import torch +from prototype_common_utils import ( + make_bounding_box, + make_detection_mask, + make_image, + make_label, + make_segmentation_mask, +) +from torch.utils.data import DataLoader from torchvision.prototype import features @@ -10,104 +20,177 @@ def test_isinstance(): ) -def test_wrapping_no_copy(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) +class TestTorchFunction: + def test_wrapping_no_copy(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) - assert label.data_ptr() == tensor.data_ptr() + assert label.data_ptr() == tensor.data_ptr() + def test_to_wrapping(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) -def test_to_wrapping(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) + label_to = label.to(torch.int32) - label_to = label.to(torch.int32) + assert type(label_to) is features.Label + assert label_to.dtype is torch.int32 + assert label_to.categories is label.categories - assert type(label_to) is features.Label - assert label_to.dtype is torch.int32 - assert label_to.categories is label.categories + def test_to_feature_reference(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]).to(torch.int32) + tensor_to = tensor.to(label) -def test_to_feature_reference(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]).to(torch.int32) + assert type(tensor_to) is torch.Tensor + assert tensor_to.dtype is torch.int32 - tensor_to = tensor.to(label) + def test_clone_wrapping(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) - assert type(tensor_to) is torch.Tensor - assert tensor_to.dtype is torch.int32 + label_clone = label.clone() + assert type(label_clone) is features.Label + assert label_clone.data_ptr() != label.data_ptr() + assert label_clone.categories is label.categories -def test_clone_wrapping(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) + def test_requires_grad__wrapping(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.float32) + label = features.Label(tensor, categories=["foo", "bar"]) - label_clone = label.clone() + assert not label.requires_grad - assert type(label_clone) is features.Label - assert label_clone.data_ptr() != label.data_ptr() - assert label_clone.categories is label.categories + label_requires_grad = label.requires_grad_(True) + assert type(label_requires_grad) is features.Label + assert label.requires_grad + assert label_requires_grad.requires_grad -def test_requires_grad__wrapping(): - tensor = torch.tensor([0, 1, 0], dtype=torch.float32) - label = features.Label(tensor, categories=["foo", "bar"]) + def test_other_op_no_wrapping(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) + + # any operation besides .to() and .clone() will do here + output = label * 2 - assert not label.requires_grad + assert type(output) is torch.Tensor - label_requires_grad = label.requires_grad_(True) + @pytest.mark.parametrize( + "op", + [ + lambda t: t.numpy(), + lambda t: t.tolist(), + lambda t: t.max(dim=-1), + ], + ) + def test_no_tensor_output_op_no_wrapping(self, op): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) - assert type(label_requires_grad) is features.Label - assert label.requires_grad - assert label_requires_grad.requires_grad + output = op(label) + assert type(output) is not features.Label -def test_other_op_no_wrapping(): + def test_inplace_op_no_wrapping(self): + tensor = torch.tensor([0, 1, 0], dtype=torch.int64) + label = features.Label(tensor, categories=["foo", "bar"]) + + output = label.add_(0) + + assert type(output) is torch.Tensor + assert type(label) is features.Label + + +def test_new_like(): tensor = torch.tensor([0, 1, 0], dtype=torch.int64) label = features.Label(tensor, categories=["foo", "bar"]) # any operation besides .to() and .clone() will do here output = label * 2 - assert type(output) is torch.Tensor + label_new = features.Label.new_like(label, output) + assert type(label_new) is features.Label + assert label_new.data_ptr() == output.data_ptr() + assert label_new.categories is label.categories -@pytest.mark.parametrize( - "op", - [ - lambda t: t.numpy(), - lambda t: t.tolist(), - lambda t: t.max(dim=-1), - ], -) -def test_no_tensor_output_op_no_wrapping(op): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) - output = op(label) +class TestVisionCollate: + def check_collation(self, dataset, expected_batch, *, collate_fn=features.vision_collate): + data_loader = DataLoader(dataset, num_workers=0, batch_size=len(dataset), collate_fn=collate_fn) - assert type(output) is not features.Label + actual_batch = list(data_loader)[0] + torch.testing.assert_close(actual_batch, expected_batch) -def test_inplace_op_no_wrapping(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) + return actual_batch - output = label.add_(0) + @pytest.mark.parametrize("with_labels", [True, False]) + def test_classification(self, with_labels): + image_size = (16, 17) + categories = ["foo", "bar", "baz"] - assert type(output) is torch.Tensor - assert type(label) is features.Label + dataset = [] + for _ in range(4): + image = make_image(size=image_size) + label = make_label(categories=categories) if with_labels else None + dataset.append((image, label)) -def test_new_like(): - tensor = torch.tensor([0, 1, 0], dtype=torch.int64) - label = features.Label(tensor, categories=["foo", "bar"]) + expected_images, expected_labels = zip(*dataset) + expected_batch = [ + features.Image.new_like(expected_images[0], torch.stack(expected_images)), + features.Label.new_like(expected_labels[0], torch.stack(expected_labels)) + if with_labels + else list(expected_labels), + ] - # any operation besides .to() and .clone() will do here - output = label * 2 + actual_batch = self.check_collation(dataset, expected_batch) - label_new = features.Label.new_like(label, output) + if with_labels: + assert actual_batch[1].categories == categories - assert type(label_new) is features.Label - assert label_new.data_ptr() == output.data_ptr() - assert label_new.categories is label.categories + def test_segmentation(self): + image_size = (16, 17) + + dataset = [] + for _ in range(4): + image = make_image(size=image_size) + mask = make_segmentation_mask(size=image_size, num_categories=10) + + dataset.append((image, mask)) + + expected_batch = [ + type(expected_features[0]).new_like(expected_features[0], torch.stack(expected_features)) + for expected_features in zip(*dataset) + ] + + self.check_collation(dataset, expected_batch) + + def test_detection(self): + dataset = [] + for _ in range(4): + image = make_image() + image_size = image.image_size + + num_objects = int(torch.randint(1, 11, ())) + + target = dict( + boxes=make_bounding_box( + extra_dims=(num_objects,), format=features.BoundingBoxFormat.XYXY, image_size=image_size + ), + labels=make_label(extra_dims=(num_objects,)), + masks=make_detection_mask(size=image_size, num_objects=num_objects), + ) + + dataset.append((image, target)) + + expected_batch = list(zip(*dataset)) + + self.check_collation( + dataset, + expected_batch, + collate_fn=functools.partial(features.vision_collate, detection_task=True), + ) diff --git a/torchvision/prototype/features/__init__.py b/torchvision/prototype/features/__init__.py index df77e8b77b3..4b1b208c3f6 100644 --- a/torchvision/prototype/features/__init__.py +++ b/torchvision/prototype/features/__init__.py @@ -13,3 +13,5 @@ ) from ._label import Label, OneHotLabel from ._mask import Mask + +from ._collate import vision_collate # usort: skip diff --git a/torchvision/prototype/features/_collate.py b/torchvision/prototype/features/_collate.py new file mode 100644 index 00000000000..f8a384e0488 --- /dev/null +++ b/torchvision/prototype/features/_collate.py @@ -0,0 +1,36 @@ +from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type, Union + +from torch.utils.data._utils.collate import collate, collate_tensor_fn, default_collate_fn_map +from torchvision.prototype.features import BoundingBox, Image, Label, Mask, OneHotLabel + + +def no_collate_fn( + batch: Sequence[Any], *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None +) -> Any: + return batch + + +def new_like_collate_fn( + batch: Sequence[Any], *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None +) -> Any: + feature = batch[0] + tensor = collate_tensor_fn(batch, collate_fn_map=collate_fn_map) + return type(feature).new_like(feature, tensor) + + +vision_default_collate_fn_map = { + (Image, Mask, Label, OneHotLabel): new_like_collate_fn, + type(None): no_collate_fn, + **default_collate_fn_map, +} + +vision_detection_collate_fn_map = { + (Image, Mask, Label, OneHotLabel, BoundingBox, type(None)): no_collate_fn, + **default_collate_fn_map, +} + + +def vision_collate(batch: Sequence[Any], *, detection_task: bool = False) -> Any: + return collate( + batch, collate_fn_map=vision_detection_collate_fn_map if detection_task else vision_default_collate_fn_map + ) From 05e6ca80d1b40dd8af1b0ac5b39f7a2d12f75a1a Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Mon, 3 Oct 2022 11:53:29 +0200 Subject: [PATCH 2/2] remove detection collation for now --- test/test_prototype_features.py | 36 +--------------------- torchvision/prototype/features/_collate.py | 15 +++------ 2 files changed, 5 insertions(+), 46 deletions(-) diff --git a/test/test_prototype_features.py b/test/test_prototype_features.py index 0a7cb7267b3..5a8c64bb73c 100644 --- a/test/test_prototype_features.py +++ b/test/test_prototype_features.py @@ -1,14 +1,6 @@ -import functools - import pytest import torch -from prototype_common_utils import ( - make_bounding_box, - make_detection_mask, - make_image, - make_label, - make_segmentation_mask, -) +from prototype_common_utils import make_image, make_label, make_segmentation_mask from torch.utils.data import DataLoader from torchvision.prototype import features @@ -168,29 +160,3 @@ def test_segmentation(self): ] self.check_collation(dataset, expected_batch) - - def test_detection(self): - dataset = [] - for _ in range(4): - image = make_image() - image_size = image.image_size - - num_objects = int(torch.randint(1, 11, ())) - - target = dict( - boxes=make_bounding_box( - extra_dims=(num_objects,), format=features.BoundingBoxFormat.XYXY, image_size=image_size - ), - labels=make_label(extra_dims=(num_objects,)), - masks=make_detection_mask(size=image_size, num_objects=num_objects), - ) - - dataset.append((image, target)) - - expected_batch = list(zip(*dataset)) - - self.check_collation( - dataset, - expected_batch, - collate_fn=functools.partial(features.vision_collate, detection_task=True), - ) diff --git a/torchvision/prototype/features/_collate.py b/torchvision/prototype/features/_collate.py index f8a384e0488..1b9816432df 100644 --- a/torchvision/prototype/features/_collate.py +++ b/torchvision/prototype/features/_collate.py @@ -1,7 +1,7 @@ from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type, Union from torch.utils.data._utils.collate import collate, collate_tensor_fn, default_collate_fn_map -from torchvision.prototype.features import BoundingBox, Image, Label, Mask, OneHotLabel +from torchvision.prototype.features import Image, Label, Mask, OneHotLabel def no_collate_fn( @@ -18,19 +18,12 @@ def new_like_collate_fn( return type(feature).new_like(feature, tensor) -vision_default_collate_fn_map = { +vision_collate_fn_map = { (Image, Mask, Label, OneHotLabel): new_like_collate_fn, type(None): no_collate_fn, **default_collate_fn_map, } -vision_detection_collate_fn_map = { - (Image, Mask, Label, OneHotLabel, BoundingBox, type(None)): no_collate_fn, - **default_collate_fn_map, -} - -def vision_collate(batch: Sequence[Any], *, detection_task: bool = False) -> Any: - return collate( - batch, collate_fn_map=vision_detection_collate_fn_map if detection_task else vision_default_collate_fn_map - ) +def vision_collate(batch: Sequence[Any]) -> Any: + return collate(batch, collate_fn_map=vision_collate_fn_map)