returnn.torch.data.extern_data

From raw dict to extern_data tensor dict.

returnn.torch.data.extern_data.extern_data_template_from_config_opts(extern_data_dict: Dict[str, Any]) → TensorDict

Parameters:

extern_data_dict – as you would specify in the config

Returns:

extern data tensor dict

returnn.torch.data.extern_data.raw_dict_to_extern_data(extern_data_raw: Dict[str, Tensor | ndarray], *, extern_data_template: TensorDict, device: str | device) → TensorDict

Parameters:

• extern_data_raw – This comes out of the DataLoader, via our collate_batch.

• extern_data_template – Specified via extern_data in the config.

• device – E.g. the GPU.

Returns:

tensor dict, like extern_data_template, but with raw tensors set to Torch tensors, on the right device.

returnn.torch.data.extern_data.raw_dict_can_split_batch(extern_data_raw: Dict[str, Tensor | ndarray], *, num_splits: int = 2) → bool

Parameters:

• extern_data_raw – This comes out of the DataLoader, via our collate_batch.

• num_splits – into how many parts

Returns:

returnn.torch.data.extern_data.raw_dict_split_batch(extern_data_raw: Dict[str, Tensor | ndarray], *, splits: Sequence[int] | int) → List[Dict[str, Tensor | ndarray]]

Parameters:

• extern_data_raw – This comes out of the DataLoader, via our collate_batch.

• splits – either the list of resulting batch dims, or the num of splits

Returns:

splitted extern_data_raw

returnn.torch.data.extern_data.get_batch_dim_from_extern_data(extern_data: TensorDict) → Dim

We expect that the batch dim is the first dim in any of the tensors. See collate_batch.

We allow that the batch dim is not necessarily the global batch_dim object. We also allow that this is not even marked as batch dim (is_batch_dim() can be False).