(tensor_list: List[Tensor])
| 308 | |
| 309 | |
| 310 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
| 311 | # TODO make this more general |
| 312 | if tensor_list[0].ndim == 3: |
| 313 | if torchvision._is_tracing(): |
| 314 | # nested_tensor_from_tensor_list() does not export well to ONNX |
| 315 | # call _onnx_nested_tensor_from_tensor_list() instead |
| 316 | return _onnx_nested_tensor_from_tensor_list(tensor_list) |
| 317 | |
| 318 | # TODO make it support different-sized images |
| 319 | max_size = _max_by_axis([list(img.shape) for img in tensor_list]) |
| 320 | # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) |
| 321 | batch_shape = [len(tensor_list)] + max_size |
| 322 | b, c, h, w = batch_shape |
| 323 | dtype = tensor_list[0].dtype |
| 324 | device = tensor_list[0].device |
| 325 | tensor = torch.zeros(batch_shape, dtype=dtype, device=device) |
| 326 | mask = torch.ones((b, h, w), dtype=torch.bool, device=device) |
| 327 | for img, pad_img, m in zip(tensor_list, tensor, mask): |
| 328 | pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) |
| 329 | m[: img.shape[1], :img.shape[2]] = False |
| 330 | else: |
| 331 | raise ValueError('not supported') |
| 332 | return NestedTensor(tensor, mask) |
| 333 | |
| 334 | |
| 335 | # _onnx_nested_tensor_from_tensor_list() is an implementation of |
no test coverage detected