Constructs a :class:`Batch` object. Constructs a batch by copying the input tensors and optionally converting them to the desired data type and storing on the specified device. Parameters ---------- tensors : TensorLike, default: None The data to construct the batch fro
(
tensors: BatchLike,
dtype: DTypeLike | None = None,
device: DeviceLike | None = None,
layout: str | None = None,
)
| 906 | |
| 907 | |
| 908 | def batch( |
| 909 | tensors: BatchLike, |
| 910 | dtype: DTypeLike | None = None, |
| 911 | device: DeviceLike | None = None, |
| 912 | layout: str | None = None, |
| 913 | ) -> Batch: |
| 914 | """Constructs a :class:`Batch` object. |
| 915 | |
| 916 | Constructs a batch by copying the input tensors and optionally converting them to the desired |
| 917 | data type and storing on the specified device. |
| 918 | |
| 919 | Parameters |
| 920 | ---------- |
| 921 | tensors : TensorLike, default: None |
| 922 | The data to construct the batch from. Can be a list of tensors, a TensorList, |
| 923 | or other supported types. |
| 924 | Supported types are: |
| 925 | |
| 926 | - a :class:`Batch` object; the batch is copied and the data is converted and moved to the |
| 927 | specified device, if necessary |
| 928 | - a list of tensor-like objects; the objects need to have matching number of dimensions, |
| 929 | data types and layouts, |
| 930 | - a tensor-like object; the outermost dimenion is interpreted as the batch dimension |
| 931 | - a dali.backend.TensorListCPU or dali.backend.TensorListGPU |
| 932 | dtype : DType, default: None |
| 933 | The desired data type of the batch. If not specified, the data type is inferred |
| 934 | from the input tensors. If specified, the input tensors are cast to the desired data type. |
| 935 | The `dtype` is required if tensors are an empty list. |
| 936 | device : Device or str, optional, default: None |
| 937 | The device on which the batch should reside (e.g., "cpu" or "gpu"). |
| 938 | If not specified, the device is inferred from the input tensors. |
| 939 | layout : str, optional, default: None |
| 940 | The layout string describing the dimensions of the batch (e.g., "HWC"). |
| 941 | If not specified, the layout is inferred from the input tensors. |
| 942 | """ |
| 943 | if isinstance(tensors, Batch): |
| 944 | b = tensors.to_device(device or tensors.device, force_copy=True) |
| 945 | if dtype is not None and b.dtype != dtype: |
| 946 | from . import cast |
| 947 | |
| 948 | b = cast(b, dtype=dtype, device=device) |
| 949 | if layout is not None and layout != b.layout: |
| 950 | from . import reshape |
| 951 | |
| 952 | b = reshape(b, layout=layout) # TODO(michalz): optimize |
| 953 | return b.evaluate() |
| 954 | else: |
| 955 | return Batch(tensors, dtype=dtype, device=device, layout=layout, copy=True) |
| 956 | |
| 957 | |
| 958 | def as_batch( |