Repository: pytorch/vision

Issue:
datapoints cannot be deepcopied
```py
from torchvision.prototype import features
from copy import copy, deepcopy
import torch

image = features.Image(torch.rand(3, 16, 16))
image_copy = copy(image)  # works
image_deepcopy = deepcopy(image)  # fails
```

```
RuntimeError: The default implementation of __deepcopy__() for non-wrapper subclasses only works for subclass types 
that implement new_empty() and for which that function returns another instance of the same subclass. 
You should either properly implement new_empty() for your subclass or override __deepcopy__() if it is intended 
behavior for new_empty() to return an instance of a different type.
```

The error message is giving us two options:

1. Override [`new_empty`](https://pytorch.org/docs/stable/generated/torch.Tensor.new_empty.html): according to the error message, this is preferred since we want to return an instance of the same type. Some of our features have internal guards for the shape like 

    https://github.com/pytorch/vision/blob/b4b246a52ebbe29bff9099c6fbabff0b82d7edd3/torchvision/prototype/features/_image.py#L49-L50

    and the plan is to expand this to all features, since the transformations rely on a certain shape. However, you can call `new_empty(size)` with an arbitrary size and the default implementation of `__deepcopy__` does so with an [empty list `[]`](https://github.com/pytorch/pytorch/blob/432c508e71111f9d5382322e0e6b1bc1c66bf0ec/torch/_tensor.py#L181). Later on, [the correct shape of the copied tensor is set](https://github.com/pytorch/pytorch/blob/432c508e71111f9d5382322e0e6b1bc1c66bf0ec/torch/_tensor.py#L191-L193), but we still need to be able to create the tensor with an incorrect shape. We could left-pad the `size` with zeros to get to the correct number of dimensions in case it is needed.

2. Override `__deepcopy__`: Maybe I'm to naive here, but a deepcopy should just be a `.clone()` of the tensor as well as a deepcopy of attributes where it is needed, e.g. `Label.categories`. Not sure if there is a downside, but that seems easier than 1. while achieving the same goal.

cc @bjuncek


Discussion:
I don't know the details of why they chose to have such a default implementation of `__deepcopy__` in option 1 but I think your argument against it makes sense to me. Option 2 does look more appealing and I agree that intuitively I would expect a clone() to be used + a deep copy of the meta-data.

@vfdev-5 thoughts?

Resolving patch:
diff --git a/torchvision/datapoints/_datapoint.py b/torchvision/datapoints/_datapoint.py
--- a/torchvision/datapoints/_datapoint.py
+++ b/torchvision/datapoints/_datapoint.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from types import ModuleType
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
 
 import PIL.Image
 import torch
@@ -36,6 +36,7 @@ def wrap_like(cls: Type[D], other: D, tensor: torch.Tensor) -> D:
     _NO_WRAPPING_EXCEPTIONS = {
         torch.Tensor.clone: lambda cls, input, output: cls.wrap_like(input, output),
         torch.Tensor.to: lambda cls, input, output: cls.wrap_like(input, output),
+        torch.Tensor.detach: lambda cls, input, output: cls.wrap_like(input, output),
         # We don't need to wrap the output of `Tensor.requires_grad_`, since it is an inplace operation and thus
         # retains the type automatically
         torch.Tensor.requires_grad_: lambda cls, input, output: output,
@@ -132,6 +133,15 @@ def dtype(self) -> _dtype:  # type: ignore[override]
         with DisableTorchFunctionSubclass():
             return super().dtype
 
+    def __deepcopy__(self: D, memo: Dict[int, Any]) -> D:
+        # We need to detach first, since a plain `Tensor.clone` will be part of the computation graph, which does
+        # *not* happen for `deepcopy(Tensor)`. A side-effect from detaching is that the `Tensor.requires_grad`
+        # attribute is cleared, so we need to refill it before we return.
+        # Note: We don't explicitly handle deep-copying of the metadata here. The only metadata we currently have is
+        # `BoundingBox.format` and `BoundingBox.spatial_size`, which are immutable and thus implicitly deep-copied by
+        # `BoundingBox.clone()`.
+        return self.detach().clone().requires_grad_(self.requires_grad)  # type: ignore[return-value]
+
     def horizontal_flip(self) -> Datapoint:
         return self
 


Test patch:
diff --git a/test/test_datapoints.py b/test/test_datapoints.py
--- a/test/test_datapoints.py
+++ b/test/test_datapoints.py
@@ -1,5 +1,8 @@
+from copy import deepcopy
+
 import pytest
 import torch
+from common_utils import assert_equal
 from PIL import Image
 
 from torchvision import datapoints
@@ -30,3 +33,154 @@ def test_bbox_instance(data, format):
     if isinstance(format, str):
         format = datapoints.BoundingBoxFormat[(format.upper())]
     assert bboxes.format == format
+
+
+@pytest.mark.parametrize(
+    ("data", "input_requires_grad", "expected_requires_grad"),
+    [
+        ([[[0.0, 1.0], [0.0, 1.0]]], None, False),
+        ([[[0.0, 1.0], [0.0, 1.0]]], False, False),
+        ([[[0.0, 1.0], [0.0, 1.0]]], True, True),
+        (torch.rand(3, 16, 16, requires_grad=False), None, False),
+        (torch.rand(3, 16, 16, requires_grad=False), False, False),
+        (torch.rand(3, 16, 16, requires_grad=False), True, True),
+        (torch.rand(3, 16, 16, requires_grad=True), None, True),
+        (torch.rand(3, 16, 16, requires_grad=True), False, False),
+        (torch.rand(3, 16, 16, requires_grad=True), True, True),
+    ],
+)
+def test_new_requires_grad(data, input_requires_grad, expected_requires_grad):
+    datapoint = datapoints.Image(data, requires_grad=input_requires_grad)
+    assert datapoint.requires_grad is expected_requires_grad
+
+
+def test_isinstance():
+    assert isinstance(datapoints.Image(torch.rand(3, 16, 16)), torch.Tensor)
+
+
+def test_wrapping_no_copy():
+    tensor = torch.rand(3, 16, 16)
+    image = datapoints.Image(tensor)
+
+    assert image.data_ptr() == tensor.data_ptr()
+
+
+def test_to_wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    image_to = image.to(torch.float64)
+
+    assert type(image_to) is datapoints.Image
+    assert image_to.dtype is torch.float64
+
+
+def test_to_datapoint_reference():
+    tensor = torch.rand((3, 16, 16), dtype=torch.float64)
+    image = datapoints.Image(tensor)
+
+    tensor_to = tensor.to(image)
+
+    assert type(tensor_to) is torch.Tensor
+    assert tensor_to.dtype is torch.float64
+
+
+def test_clone_wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    image_clone = image.clone()
+
+    assert type(image_clone) is datapoints.Image
+    assert image_clone.data_ptr() != image.data_ptr()
+
+
+def test_requires_grad__wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    assert not image.requires_grad
+
+    image_requires_grad = image.requires_grad_(True)
+
+    assert type(image_requires_grad) is datapoints.Image
+    assert image.requires_grad
+    assert image_requires_grad.requires_grad
+
+
+def test_detach_wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16), requires_grad=True)
+
+    image_detached = image.detach()
+
+    assert type(image_detached) is datapoints.Image
+
+
+def test_other_op_no_wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    # any operation besides the ones listed in `Datapoint._NO_WRAPPING_EXCEPTIONS` will do here
+    output = image * 2
+
+    assert type(output) is torch.Tensor
+
+
+@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):
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    output = op(image)
+
+    assert type(output) is not datapoints.Image
+
+
+def test_inplace_op_no_wrapping():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    output = image.add_(0)
+
+    assert type(output) is torch.Tensor
+    assert type(image) is datapoints.Image
+
+
+def test_wrap_like():
+    image = datapoints.Image(torch.rand(3, 16, 16))
+
+    # any operation besides the ones listed in `Datapoint._NO_WRAPPING_EXCEPTIONS` will do here
+    output = image * 2
+
+    image_new = datapoints.Image.wrap_like(image, output)
+
+    assert type(image_new) is datapoints.Image
+    assert image_new.data_ptr() == output.data_ptr()
+
+
+@pytest.mark.parametrize(
+    "datapoint",
+    [
+        datapoints.Image(torch.rand(3, 16, 16)),
+        datapoints.Video(torch.rand(2, 3, 16, 16)),
+        datapoints.BoundingBox([0.0, 1.0, 2.0, 3.0], format=datapoints.BoundingBoxFormat.XYXY, spatial_size=(10, 10)),
+        datapoints.Mask(torch.randint(0, 256, (16, 16), dtype=torch.uint8)),
+    ],
+)
+@pytest.mark.parametrize("requires_grad", [False, True])
+def test_deepcopy(datapoint, requires_grad):
+    if requires_grad and not datapoint.dtype.is_floating_point:
+        return
+
+    datapoint.requires_grad_(requires_grad)
+
+    datapoint_deepcopied = deepcopy(datapoint)
+
+    assert datapoint_deepcopied is not datapoint
+    assert datapoint_deepcopied.data_ptr() != datapoint.data_ptr()
+    assert_equal(datapoint_deepcopied, datapoint)
+
+    assert type(datapoint_deepcopied) is type(datapoint)
+    assert datapoint_deepcopied.requires_grad is requires_grad
+    assert datapoint_deepcopied.is_leaf
diff --git a/test/test_prototype_datapoints.py b/test/test_prototype_datapoints.py
deleted file mode 100644
--- a/test/test_prototype_datapoints.py
+++ /dev/null
@@ -1,133 +0,0 @@
-import pytest
-import torch
-
-from torchvision.prototype import datapoints as proto_datapoints
-
-
-@pytest.mark.parametrize(
-    ("data", "input_requires_grad", "expected_requires_grad"),
-    [
-        ([0.0], None, False),
-        ([0.0], False, False),
-        ([0.0], True, True),
-        (torch.tensor([0.0], requires_grad=False), None, False),
-        (torch.tensor([0.0], requires_grad=False), False, False),
-        (torch.tensor([0.0], requires_grad=False), True, True),
-        (torch.tensor([0.0], requires_grad=True), None, True),
-        (torch.tensor([0.0], requires_grad=True), False, False),
-        (torch.tensor([0.0], requires_grad=True), True, True),
-    ],
-)
-def test_new_requires_grad(data, input_requires_grad, expected_requires_grad):
-    datapoint = proto_datapoints.Label(data, requires_grad=input_requires_grad)
-    assert datapoint.requires_grad is expected_requires_grad
-
-
-def test_isinstance():
-    assert isinstance(
-        proto_datapoints.Label([0, 1, 0], categories=["foo", "bar"]),
-        torch.Tensor,
-    )
-
-
-def test_wrapping_no_copy():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    assert label.data_ptr() == tensor.data_ptr()
-
-
-def test_to_wrapping():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    label_to = label.to(torch.int32)
-
-    assert type(label_to) is proto_datapoints.Label
-    assert label_to.dtype is torch.int32
-    assert label_to.categories is label.categories
-
-
-def test_to_datapoint_reference():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"]).to(torch.int32)
-
-    tensor_to = tensor.to(label)
-
-    assert type(tensor_to) is torch.Tensor
-    assert tensor_to.dtype is torch.int32
-
-
-def test_clone_wrapping():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    label_clone = label.clone()
-
-    assert type(label_clone) is proto_datapoints.Label
-    assert label_clone.data_ptr() != label.data_ptr()
-    assert label_clone.categories is label.categories
-
-
-def test_requires_grad__wrapping():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.float32)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    assert not label.requires_grad
-
-    label_requires_grad = label.requires_grad_(True)
-
-    assert type(label_requires_grad) is proto_datapoints.Label
-    assert label.requires_grad
-    assert label_requires_grad.requires_grad
-
-
-def test_other_op_no_wrapping():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    # any operation besides .to() and .clone() will do here
-    output = label * 2
-
-    assert type(output) is torch.Tensor
-
-
-@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 = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    output = op(label)
-
-    assert type(output) is not proto_datapoints.Label
-
-
-def test_inplace_op_no_wrapping():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    output = label.add_(0)
-
-    assert type(output) is torch.Tensor
-    assert type(label) is proto_datapoints.Label
-
-
-def test_wrap_like():
-    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
-    label = proto_datapoints.Label(tensor, categories=["foo", "bar"])
-
-    # any operation besides .to() and .clone() will do here
-    output = label * 2
-
-    label_new = proto_datapoints.Label.wrap_like(label, output)
-
-    assert type(label_new) is proto_datapoints.Label
-    assert label_new.data_ptr() == output.data_ptr()
-    assert label_new.categories is label.categories