1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
|
from copy import deepcopy
import pytest
import torch
from common_utils import assert_equal, make_bounding_boxes, make_image, make_segmentation_mask, make_video
from PIL import Image
from torchvision import tv_tensors
@pytest.fixture(autouse=True)
def restore_tensor_return_type():
# This is for security, as we should already be restoring the default manually in each test anyway
# (at least at the time of writing...)
yield
tv_tensors.set_return_type("Tensor")
@pytest.mark.parametrize("data", [torch.rand(3, 32, 32), Image.new("RGB", (32, 32), color=123)])
def test_image_instance(data):
image = tv_tensors.Image(data)
assert isinstance(image, torch.Tensor)
assert image.ndim == 3 and image.shape[0] == 3
@pytest.mark.parametrize("data", [torch.randint(0, 10, size=(1, 32, 32)), Image.new("L", (32, 32), color=2)])
def test_mask_instance(data):
mask = tv_tensors.Mask(data)
assert isinstance(mask, torch.Tensor)
assert mask.ndim == 3 and mask.shape[0] == 1
@pytest.mark.parametrize("data", [torch.randint(0, 32, size=(5, 4)), [[0, 0, 5, 5], [2, 2, 7, 7]], [1, 2, 3, 4]])
@pytest.mark.parametrize(
"format", ["XYXY", "CXCYWH", tv_tensors.BoundingBoxFormat.XYXY, tv_tensors.BoundingBoxFormat.XYWH]
)
def test_bbox_instance(data, format):
bboxes = tv_tensors.BoundingBoxes(data, format=format, canvas_size=(32, 32))
assert isinstance(bboxes, torch.Tensor)
assert bboxes.ndim == 2 and bboxes.shape[1] == 4
if isinstance(format, str):
format = tv_tensors.BoundingBoxFormat[(format.upper())]
assert bboxes.format == format
def test_bbox_dim_error():
data_3d = [[[1, 2, 3, 4]]]
with pytest.raises(ValueError, match="Expected a 1D or 2D tensor, got 3D"):
tv_tensors.BoundingBoxes(data_3d, format="XYXY", canvas_size=(32, 32))
@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):
tv_tensor = tv_tensors.Image(data, requires_grad=input_requires_grad)
assert tv_tensor.requires_grad is expected_requires_grad
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
def test_isinstance(make_input):
assert isinstance(make_input(), torch.Tensor)
def test_wrapping_no_copy():
tensor = torch.rand(3, 16, 16)
image = tv_tensors.Image(tensor)
assert image.data_ptr() == tensor.data_ptr()
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
def test_to_wrapping(make_input):
dp = make_input()
dp_to = dp.to(torch.float64)
assert type(dp_to) is type(dp)
assert dp_to.dtype is torch.float64
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_to_tv_tensor_reference(make_input, return_type):
tensor = torch.rand((3, 16, 16), dtype=torch.float64)
dp = make_input()
with tv_tensors.set_return_type(return_type):
tensor_to = tensor.to(dp)
assert type(tensor_to) is (type(dp) if return_type == "TVTensor" else torch.Tensor)
assert tensor_to.dtype is dp.dtype
assert type(tensor) is torch.Tensor
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_clone_wrapping(make_input, return_type):
dp = make_input()
with tv_tensors.set_return_type(return_type):
dp_clone = dp.clone()
assert type(dp_clone) is type(dp)
assert dp_clone.data_ptr() != dp.data_ptr()
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_requires_grad__wrapping(make_input, return_type):
dp = make_input(dtype=torch.float)
assert not dp.requires_grad
with tv_tensors.set_return_type(return_type):
dp_requires_grad = dp.requires_grad_(True)
assert type(dp_requires_grad) is type(dp)
assert dp.requires_grad
assert dp_requires_grad.requires_grad
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_detach_wrapping(make_input, return_type):
dp = make_input(dtype=torch.float).requires_grad_(True)
with tv_tensors.set_return_type(return_type):
dp_detached = dp.detach()
assert type(dp_detached) is type(dp)
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_force_subclass_with_metadata(return_type):
# Sanity checks for the ops in _FORCE_TORCHFUNCTION_SUBCLASS and tv_tensors with metadata
# Largely the same as above, we additionally check that the metadata is preserved
format, canvas_size = "XYXY", (32, 32)
bbox = tv_tensors.BoundingBoxes([[0, 0, 5, 5], [2, 2, 7, 7]], format=format, canvas_size=canvas_size)
tv_tensors.set_return_type(return_type)
bbox = bbox.clone()
if return_type == "TVTensor":
assert bbox.format, bbox.canvas_size == (format, canvas_size)
bbox = bbox.to(torch.float64)
if return_type == "TVTensor":
assert bbox.format, bbox.canvas_size == (format, canvas_size)
bbox = bbox.detach()
if return_type == "TVTensor":
assert bbox.format, bbox.canvas_size == (format, canvas_size)
assert not bbox.requires_grad
bbox.requires_grad_(True)
if return_type == "TVTensor":
assert bbox.format, bbox.canvas_size == (format, canvas_size)
assert bbox.requires_grad
tv_tensors.set_return_type("tensor")
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_other_op_no_wrapping(make_input, return_type):
dp = make_input()
with tv_tensors.set_return_type(return_type):
# any operation besides the ones listed in _FORCE_TORCHFUNCTION_SUBCLASS will do here
output = dp * 2
assert type(output) is (type(dp) if return_type == "TVTensor" else torch.Tensor)
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@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(make_input, op):
dp = make_input()
output = op(dp)
assert type(output) is not type(dp)
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
def test_inplace_op_no_wrapping(make_input, return_type):
dp = make_input()
original_type = type(dp)
with tv_tensors.set_return_type(return_type):
output = dp.add_(0)
assert type(output) is (type(dp) if return_type == "TVTensor" else torch.Tensor)
assert type(dp) is original_type
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
def test_wrap(make_input):
dp = make_input()
# any operation besides the ones listed in _FORCE_TORCHFUNCTION_SUBCLASS will do here
output = dp * 2
dp_new = tv_tensors.wrap(output, like=dp)
assert type(dp_new) is type(dp)
assert dp_new.data_ptr() == output.data_ptr()
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("requires_grad", [False, True])
def test_deepcopy(make_input, requires_grad):
dp = make_input(dtype=torch.float)
dp.requires_grad_(requires_grad)
dp_deepcopied = deepcopy(dp)
assert dp_deepcopied is not dp
assert dp_deepcopied.data_ptr() != dp.data_ptr()
assert_equal(dp_deepcopied, dp)
assert type(dp_deepcopied) is type(dp)
assert dp_deepcopied.requires_grad is requires_grad
@pytest.mark.parametrize("make_input", [make_image, make_bounding_boxes, make_segmentation_mask, make_video])
@pytest.mark.parametrize("return_type", ["Tensor", "TVTensor"])
@pytest.mark.parametrize(
"op",
(
lambda dp: dp + torch.rand(*dp.shape),
lambda dp: torch.rand(*dp.shape) + dp,
lambda dp: dp * torch.rand(*dp.shape),
lambda dp: torch.rand(*dp.shape) * dp,
lambda dp: dp + 3,
lambda dp: 3 + dp,
lambda dp: dp + dp,
lambda dp: dp.sum(),
lambda dp: dp.reshape(-1),
lambda dp: dp.int(),
lambda dp: torch.stack([dp, dp]),
lambda dp: torch.chunk(dp, 2)[0],
lambda dp: torch.unbind(dp)[0],
),
)
def test_usual_operations(make_input, return_type, op):
dp = make_input()
with tv_tensors.set_return_type(return_type):
out = op(dp)
assert type(out) is (type(dp) if return_type == "TVTensor" else torch.Tensor)
if isinstance(dp, tv_tensors.BoundingBoxes) and return_type == "TVTensor":
assert hasattr(out, "format")
assert hasattr(out, "canvas_size")
def test_subclasses():
img = make_image()
masks = make_segmentation_mask()
with pytest.raises(TypeError, match="unsupported operand"):
img + masks
def test_set_return_type():
img = make_image()
assert type(img + 3) is torch.Tensor
with tv_tensors.set_return_type("TVTensor"):
assert type(img + 3) is tv_tensors.Image
assert type(img + 3) is torch.Tensor
tv_tensors.set_return_type("TVTensor")
assert type(img + 3) is tv_tensors.Image
with tv_tensors.set_return_type("tensor"):
assert type(img + 3) is torch.Tensor
with tv_tensors.set_return_type("TVTensor"):
assert type(img + 3) is tv_tensors.Image
tv_tensors.set_return_type("tensor")
assert type(img + 3) is torch.Tensor
assert type(img + 3) is torch.Tensor
# Exiting a context manager will restore the return type as it was prior to entering it,
# regardless of whether the "global" tv_tensors.set_return_type() was called within the context manager.
assert type(img + 3) is tv_tensors.Image
tv_tensors.set_return_type("tensor")
def test_return_type_input():
img = make_image()
# Case-insensitive
with tv_tensors.set_return_type("tvtensor"):
assert type(img + 3) is tv_tensors.Image
with pytest.raises(ValueError, match="return_type must be"):
tv_tensors.set_return_type("typo")
tv_tensors.set_return_type("tensor")
|
