(tensor_list: List[Tensor])
| 443 | |
| 444 | |
| 445 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
| 446 | # TODO make this more general |
| 447 | if tensor_list[0].ndim == 3: |
| 448 | if torchvision._is_tracing(): |
| 449 | # nested_tensor_from_tensor_list() does not export well to ONNX |
| 450 | # call _onnx_nested_tensor_from_tensor_list() instead |
| 451 | return _onnx_nested_tensor_from_tensor_list(tensor_list) |
| 452 | |
| 453 | # TODO make it support different-sized images |
| 454 | max_size = _max_by_axis([list(img.shape) for img in tensor_list]) |
| 455 | # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) |
| 456 | batch_shape = [len(tensor_list)] + max_size |
| 457 | b, c, h, w = batch_shape |
| 458 | dtype = tensor_list[0].dtype |
| 459 | device = tensor_list[0].device |
| 460 | tensor = torch.zeros(batch_shape, dtype=dtype, device=device) |
| 461 | mask = torch.ones((b, h, w), dtype=torch.bool, device=device) |
| 462 | for img, pad_img, m in zip(tensor_list, tensor, mask): |
| 463 | pad_img[:img.shape[0], :img.shape[1], :img.shape[2]].copy_(img) |
| 464 | m[:img.shape[1], :img.shape[2]] = False |
| 465 | else: |
| 466 | raise ValueError('not supported') |
| 467 | return NestedTensor(tensor, mask) |
| 468 | |
| 469 | |
| 470 | @torch.jit.unused |
no test coverage detected