(tensor_list: List[Tensor])
| 373 | |
| 374 | |
| 375 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
| 376 | # TODO make this more general |
| 377 | if tensor_list[0].ndim == 3: |
| 378 | if torchvision._is_tracing(): |
| 379 | # nested_tensor_from_tensor_list() does not export well to ONNX |
| 380 | # call _onnx_nested_tensor_from_tensor_list() instead |
| 381 | return _onnx_nested_tensor_from_tensor_list(tensor_list) |
| 382 | |
| 383 | # TODO make it support different-sized images |
| 384 | max_size = _max_by_axis([list(img.shape) for img in tensor_list]) |
| 385 | # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) |
| 386 | batch_shape = [len(tensor_list)] + max_size |
| 387 | b, c, h, w = batch_shape |
| 388 | dtype = tensor_list[0].dtype |
| 389 | device = tensor_list[0].device |
| 390 | tensor = torch.zeros(batch_shape, dtype=dtype, device=device) |
| 391 | mask = torch.ones((b, h, w), dtype=torch.bool, device=device) |
| 392 | for img, pad_img, m in zip(tensor_list, tensor, mask): |
| 393 | pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) |
| 394 | m[: img.shape[1], :img.shape[2]] = False |
| 395 | else: |
| 396 | raise ValueError('not supported') |
| 397 | return NestedTensor(tensor, mask) |
| 398 | |
| 399 | |
| 400 | # _onnx_nested_tensor_from_tensor_list() is an implementation of |
no test coverage detected