TensorDictBase.to_mds#

Writes the content of a TensorCollection to a streaming dataset.

Keyword Arguments:

out (str | Tuple[str, str]) –
Output dataset directory to save shard files.
1. If out is a local directory, shard files are saved locally.
2. If out is a remote directory, a local temporary directory is created to
  cache the shard files and then the shard files are uploaded to a remote location. At the end, the temp directory is deleted once shards are uploaded.
3. If out is a tuple of (local_dir, remote_dir), shard files are saved in the
  local_dir and also uploaded to a remote location.
columns (dict[str, str]) – an optional dict of columns. Will be automatically inferred from the tensor collection if not provided.
keep_local (bool) – If the dataset is uploaded, whether to keep the local dataset directory or remove it after uploading. Defaults to False.
compression (str, optional) – Optional compression or compression:level. Defaults to None.
hashes (List[str], optional) – Optional list of hash algorithms to apply to shard files. Defaults to None.
size_limit (Union[int, str], optional) – Optional shard size limit, after which point to start a new shard. If None, puts everything in one shard. Can specify bytes human-readable format as well, for example "100kb" for 100 kilobyte (100*1024) and so on. Defaults to 1 << 26. Ignored if writer is passed.
writer (MDSWriter, optional) – the writer to use. If None, it is created from out and the other keyword arguments.
**kwargs (Any) – Additional settings for the Writer.

Note

The MDSWriter has limited support for nested dictionaries. The proper way to handle nested tensordicts is to use the flatten_keys() method before writing, and unflatten_keys() method after reading.

Warning

This method requires mosaicml-streaming to be installed.

Warning

For non-tensor data, the type of the data must be fixed. The way tensordict recovers the data type is by looking at the first element of the list. If it’s None an error will be thrown. Otherwise all the data in the list must have the same type (or be None if missing).

See also

See the Mosaic streaming library API at https://docs.mosaicml.com/projects/streaming

The following example writes a dataset and loads it in a PyTorch dataloader. It needs mosaicml-streaming, so it is not run when the documentation is tested.

Examples

>>> import tempfile
>>> from typing import Any
>>> from tensordict import TensorDict, LazyStackedTensorDict
>>> import torch
>>>
>>>
>>> td = LazyStackedTensorDict(
...     TensorDict(a=0, b=1, c=torch.randn(2), d="a string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="another string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="yet another string"),
... )
>>>
>>> with tempfile.TemporaryDirectory() as tempdir:  
...     # Create a dataset on one process / thread / node...
...     td.to_mds(out=tempdir)
...
...     # Create a dataloader
...     from streaming import StreamingDataset
...
...     # Load the dataset on another thread / process / node...
...     dataset = StreamingDataset(local=tempdir, remote=None, batch_size=2)
...
...     # Use the class `from_list` method as a collate_fn
...     dl = torch.utils.data.DataLoader(dataset=dataset, batch_size=2, collate_fn=LazyStackedTensorDict.from_list)
...     for batch in dl:
...         print("batch", batch)

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.to_mds#

Docs

Tutorials

Resources