Repository: kornia/kornia
Pull request: fix TensorWrapper serialization
#### Changes

Fixes #2126 

- add also `TensorWarpper.unwrap` method

#### Type of change
<!-- Please delete options that are not relevant. -->
- [ ] 📚  Documentation Update
- [x] 🧪 Tests Cases
- [x] 🐞 Bug fix (non-breaking change which fixes an issue)
- [ ] 🔬 New feature (non-breaking change which adds functionality)
- [ ] 🚨 Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] 📝 This change requires a documentation update


#### Checklist

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] Did you update CHANGELOG in case of a major change?


Problem statement:
Round trip `torch.save` / `torch.load` not working for `Hyperplane`
### Describe the bug

Crashing when trying to save and reload a `Hyperplane` using `torch`:

```
Traceback (most recent call last):
  File "save_hyperplane.py", line 16, in <module>
    plane = torch.load("./saved_plane.pt")
  File "/home/kyle/venv/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/kyle/venv/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/kyle/venv/lib/python3.8/site-packages/kornia/core/tensor_wrapper.py", line 48, in __getattr__
    self.used_attrs.add(name)
  File "/home/kyle/venv/lib/python3.8/site-packages/kornia/core/tensor_wrapper.py", line 48, in __getattr__
    self.used_attrs.add(name)
  File "/home/kyle/venv/lib/python3.8/site-packages/kornia/core/tensor_wrapper.py", line 48, in __getattr__
    self.used_attrs.add(name)
  [Previous line repeated 993 more times]
RecursionError: maximum recursion depth exceeded
```

### Reproduction steps

```bash
$ python3 -m venv venv
$ source venv/bin/activate
(venv) $ pip install kornia
(venv) $ pip install numpy
(venv) $ python save_hyperplane.py
```

Content of save_hyperplane.py

```
from kornia.geometry.plane import Hyperplane
from kornia.geometry.vector import Vec3, Scalar
import torch

plane = Hyperplane(
    Vec3(torch.tensor([0,0,1])), Scalar(torch.tensor(0.5))
)

torch.save(plane, "./saved_plane.pt")
plane = torch.load("./saved_plane.pt")
```

### Expected behavior

Hyperplane is able to be saved and reloaded with `torch.save` & `torch.load` 

### Environment

```shell
Collecting environment information...
PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080 Laptop GPU
Nvidia driver version: 470.141.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==1.13.1
[conda] Could not collect
```

### Additional context

_No response_

Hints from issue discussion:
@Hackerman342 will take a look soon. This i guess is after introducing Vector3 and Scalar types

Commit message: fix TensorWrapper serialization

Patch to review:
diff --git a/kornia/core/tensor_wrapper.py b/kornia/core/tensor_wrapper.py
index 029f403cfa..51133714ff 100644
--- a/kornia/core/tensor_wrapper.py
+++ b/kornia/core/tensor_wrapper.py
@@ -35,14 +35,21 @@ def __init__(self, data: Tensor) -> None:
         self.__dict__["used_attrs"] = set()
         self.__dict__["used_calls"] = set()
 
-    @property
-    def data(self) -> Tensor:
-        return self._data
+    def unwrap(self):
+        return unwrap(self)
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
 
     def __repr__(self) -> str:
-        return f"{self.data}"
+        return f"{self._data}"
 
     def __getattr__(self, name):
+        if name == "data":
+            return self._data
         if name in self.__dict__:
             return self.__dict__[name]
         self.used_attrs.add(name)
diff --git a/test/core/test_tensor_wrapper.py b/test/core/test_tensor_wrapper.py
index 9b1dfa55dd..82d72b7253 100644
--- a/test/core/test_tensor_wrapper.py
+++ b/test/core/test_tensor_wrapper.py
@@ -1,8 +1,10 @@
+import tempfile
+
 import pytest
 import torch
 
 from kornia.core import Tensor
-from kornia.core.tensor_wrapper import TensorWrapper, unwrap, wrap
+from kornia.core.tensor_wrapper import TensorWrapper, wrap
 from kornia.testing import BaseTester
 
 
@@ -15,15 +17,24 @@ def test_smoke(self, device, dtype):
         assert tensor.shape == (1, 2, 3, 4)
         assert tensor.device == device
         assert tensor.dtype == dtype
-        self.assert_close(data, tensor.data)
-        self.assert_close(data, unwrap(tensor))
+        self.assert_close(data, tensor.unwrap())
+
+    def test_serialization(self, device, dtype):
+        data = torch.rand(1, 2, 3, 4, device=device, dtype=dtype)
+        tensor: TensorWrapper = wrap(data, TensorWrapper)
+
+        with tempfile.NamedTemporaryFile() as tmp:
+            file_path = tmp.name + ".pt"
+            torch.save(tensor, file_path)
+            loaded_tensor: TensorWrapper = torch.load(file_path)
+            self.assert_close(loaded_tensor.unwrap(), tensor.unwrap())
 
     def test_wrap_list(self, device, dtype):
         data_list = [torch.rand(2, device=device, dtype=dtype), torch.rand(3, device=device, dtype=dtype)]
         tensor_list = wrap(data_list, TensorWrapper)
         assert isinstance(tensor_list, list)
         assert len(tensor_list) == 2
-        tensor_list_data = unwrap(tensor_list)
+        tensor_list_data = tensor_list.unwrap()
         assert len(tensor_list_data) == 2
         self.assert_close(tensor_list_data[0], data_list[0])
         self.assert_close(tensor_list_data[1], data_list[1])
diff --git a/test/geometry/test_plane.py b/test/geometry/test_plane.py
index f88a3fb8e7..060909879d 100644
--- a/test/geometry/test_plane.py
+++ b/test/geometry/test_plane.py
@@ -1,3 +1,5 @@
+import tempfile
+
 import pytest
 import torch
 
@@ -49,6 +51,16 @@ def test_smoke(self, device, dtype, shape):
         assert pl0.normal.shape == shape or () + (3,)
         assert pl0.offset.shape == (shape + () if shape is not None else ())
 
+    def test_serialization(self, device, dtype):
+        p = Vector3.random((), device, dtype)
+        n = Vector3.random((), device, dtype).normalized()
+        plane = Hyperplane.from_vector(n, p)
+        with tempfile.NamedTemporaryFile() as tmp:
+            file_path = tmp.name + ".pt"
+            torch.save(plane, file_path)
+            loaded_plane = torch.load(file_path)
+            self.assert_close(plane.normal.unwrap(), loaded_plane.normal.unwrap())
+
     # TODO: implement `Vector2`
     # @pytest.mark.parametrize("batch_size", [1, 2])
     # def test_through_two(self, device, dtype, batch_size):


Code review comments:
diff_hunk: @@ -35,14 +35,21 @@ def __init__(self, data: Tensor) -> None:
         self.__dict__["used_attrs"] = set()
         self.__dict__["used_calls"] = set()
 
-    @property
-    def data(self) -> Tensor:
-        return self._data
+    def unwrap(self):
+        return unwrap(self)
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
 
     def __repr__(self) -> str:
-        return f"{self.data}"
+        return f"{self._data}"
 
     def __getattr__(self, name):
+        if name == "data":
+            return self._data
         if name in self.__dict__:
path: kornia/core/tensor_wrapper.py
text: @user1:
```suggestion
            return self._data
        elif name in self.__dict__:
```

diff_hunk: @@ -15,15 +17,24 @@ def test_smoke(self, device, dtype):
         assert tensor.shape == (1, 2, 3, 4)
         assert tensor.device == device
         assert tensor.dtype == dtype
-        self.assert_close(data, tensor.data)
-        self.assert_close(data, unwrap(tensor))
+        self.assert_close(data, tensor.unwrap())
+
+    def test_serialization(self, device, dtype):
+        data = torch.rand(1, 2, 3, 4, device=device, dtype=dtype)
+        tensor: TensorWrapper = wrap(data, TensorWrapper)
+
+        with tempfile.NamedTemporaryFile() as tmp:
+            file_path = tmp.name + ".pt"
+            torch.save(tensor, file_path)
+            loaded_tensor: TensorWrapper = torch.load(file_path)
+            self.assert_close(loaded_tensor.unwrap(), tensor.unwrap())
 
     def test_wrap_list(self, device, dtype):
         data_list = [torch.rand(2, device=device, dtype=dtype), torch.rand(3, device=device, dtype=dtype)]
         tensor_list = wrap(data_list, TensorWrapper)
         assert isinstance(tensor_list, list)
         assert len(tensor_list) == 2
-        tensor_list_data = unwrap(tensor_list)
+        tensor_list_data = tensor_list.unwrap()
path: test/core/test_tensor_wrapper.py
text: @user1:
not sure of the idea here

@author:
It will return the data structure as plane tensor. For now is a requirement of assert allclose, and some functions I detected that torch doesn’t like wrappers of tensors 

@user1:
But should this be a list of tensors, or a tensor itself? -- asking because this was asserted to be a list, and a list will not have this method

@author:
oh you are right. I think we can decide the behavior here but the general idea for the wrapper is that we cast Any to have a Tensor behavior. Open to ideas here

@user1:
Well, just if we try to lead internally with having a sequence of tensor which I think probably will be tricky

otherwise, we can do like is with tensors, something like:
```python
data_list = [
    5, 6.0, [7.0], # i think this should work too, no?
   torch.rand(2, device=device, dtype=dtype), torch.rand(3, device=device, dtype=dtype)
]

a = wrap(data_list, TensorWrapper)
assert isinstance(a, Sequence) and all(isinstance(x, TensorWrapper) for x in a)

b = []
for i in range(len(a)):
    # do something to each tensor
    b.append(a[i].unwrap())

assert  isinstance(b, list) and all(isinstance(x, Tensor) for x in a)
```

I think this second assert will be true if each item be a `TensorWrapper` too, because inheritance

diff_hunk: @@ -1,3 +1,5 @@
+import tempfile
path: test/geometry/test_plane.py
text: @user1:
We should avoid this in tests in favor of `tmp_path` / `tmp_path_factory` from pytest

diff_hunk: @@ -49,6 +51,16 @@ def test_smoke(self, device, dtype, shape):
         assert pl0.normal.shape == shape or () + (3,)
         assert pl0.offset.shape == (shape + () if shape is not None else ())
 
+    def test_serialization(self, device, dtype):
path: test/geometry/test_plane.py
text: @user1:
```suggestion
    def test_serialization(self, device, dtype, tmp_path):
```

from [pytest docs](https://docs.pytest.org/en/7.1.x/how-to/tmp_path.html#the-tmp-path-factory-fixture)

diff_hunk: @@ -15,15 +17,24 @@ def test_smoke(self, device, dtype):
         assert tensor.shape == (1, 2, 3, 4)
         assert tensor.device == device
         assert tensor.dtype == dtype
-        self.assert_close(data, tensor.data)
-        self.assert_close(data, unwrap(tensor))
+        self.assert_close(data, tensor.unwrap())
+
+    def test_serialization(self, device, dtype):
path: test/core/test_tensor_wrapper.py
text: @user1:
```suggestion
    def test_serialization(self, device, dtype, tmp_path):
```

diff_hunk: @@ -15,15 +17,24 @@ def test_smoke(self, device, dtype):
         assert tensor.shape == (1, 2, 3, 4)
         assert tensor.device == device
         assert tensor.dtype == dtype
-        self.assert_close(data, tensor.data)
-        self.assert_close(data, unwrap(tensor))
+        self.assert_close(data, tensor.unwrap())
+
+    def test_serialization(self, device, dtype):
+        data = torch.rand(1, 2, 3, 4, device=device, dtype=dtype)
+        tensor: TensorWrapper = wrap(data, TensorWrapper)
+
+        with tempfile.NamedTemporaryFile() as tmp:
+            file_path = tmp.name + ".pt"
+            torch.save(tensor, file_path)
+            loaded_tensor: TensorWrapper = torch.load(file_path)
+            self.assert_close(loaded_tensor.unwrap(), tensor.unwrap())
path: test/core/test_tensor_wrapper.py
text: @user1:
```suggestion
        file_path = tmp_path.mktemp("data") / "tensor.pt"
        torch.save(tensor, file_path)
        
        assert file_path.is_file()
        
        loaded_tensor: TensorWrapper = torch.load(file_path)
        self.assert_close(loaded_tensor.unwrap(), tensor.unwrap())
```

diff_hunk: @@ -49,6 +51,16 @@ def test_smoke(self, device, dtype, shape):
         assert pl0.normal.shape == shape or () + (3,)
         assert pl0.offset.shape == (shape + () if shape is not None else ())
 
+    def test_serialization(self, device, dtype):
+        p = Vector3.random((), device, dtype)
+        n = Vector3.random((), device, dtype).normalized()
+        plane = Hyperplane.from_vector(n, p)
+        with tempfile.NamedTemporaryFile() as tmp:
+            file_path = tmp.name + ".pt"
+            torch.save(plane, file_path)
+            loaded_plane = torch.load(file_path)
+            self.assert_close(plane.normal.unwrap(), loaded_plane.normal.unwrap())
path: test/geometry/test_plane.py
text: @user1:
```suggestion
        file_path = tmp_path.mktemp("data") / "hyperplane.pt"
        torch.save(plane, file_path)
        
        assert file_path.is_file()
        
        loaded_plane = torch.load(file_path)
        self.assert_close(plane.normal.unwrap(), loaded_plane.normal.unwrap())
```