MemoryMappedTensor

class tensordict.MemoryMappedTensor(source, *, dtype=None, shape=None, index=None, device=None, handler=None, filename=None)

A Memory-mapped Tensor.

Supports filenames or file handlers.

The main advantage of MemoryMappedTensor resides in its serialization methods, which ensure that the tensor is passed through queues or RPC remote calls without any copy.

Note

When used within RPC settings, the filepath should be accessible to both nodes. If it isn’t the behaviour of passing a MemoryMappedTensor from one worker to another is undefined.

MemoryMappedTensor supports multiple construction methods.

Examples

>>> # from an existing tensor
>>> tensor = torch.randn(3)
>>> with tempfile.NamedTemporaryFile() as file:
...     memmap_tensor = MemoryMappedTensor.from_tensor(tensor, filename=file.name)
...     assert memmap_tensor.filename is not None
>>> # if no filename is passed, a handler is used
>>> tensor = torch.randn(3)
>>> memmap_tensor = MemoryMappedTensor.from_tensor(tensor, filename=file.name)
>>> assert memmap_tensor.filename is None
>>> # one can create an empty tensor too
>>> with tempfile.NamedTemporaryFile() as file:
...     memmap_tensor_empty = MemoryMappedTensor.empty_like(tensor, filename=file.name)
>>> with tempfile.NamedTemporaryFile() as file:
...     memmap_tensor_zero = MemoryMappedTensor.zeros_like(tensor, filename=file.name)
>>> with tempfile.NamedTemporaryFile() as file:
...     memmap_tensor = MemoryMappedTensor.ones_like(tensor, filename=file.name)

chunk(chunks, dim=0) → List of Tensors

See torch.chunk()

classmethod empty(*size, dtype=None, device=None, filename=None)

classmethod empty(shape, *, dtype=None, device=None, filename=None)

Creates a tensor with empty content, specific shape, dtype and filename.

Parameters:

shape (integers or torch.Size) – the shape of the tensor.

Keyword Arguments:

• dtype (torch.dtype) – the dtype of the tensor.

• device (torch.device) – the device of the tensor. Only None and “cpu” are accepted, any other device will raise an exception.

• filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

• existsok (bool, optional) – whether it is ok to overwrite an existing file. Defaults to False.

classmethod empty_like(input, *, filename=None)

Creates a tensor with no content but the same shape and dtype as the input tensor.

Parameters:

input (torch.Tensor) – the tensor to use as an example.

Keyword Arguments:

filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

classmethod empty_nested(*args, **kwargs)

Creates a tensor with empty content, specific shape, dtype and filename.

Parameters:

shape (nested_shape) – the shapes of the tensors.

Keyword Arguments:

• dtype (torch.dtype) – the dtype of the tensor.

• device (torch.device) – the device of the tensor. Only None and “cpu” are accepted, any other device will raise an exception.

• filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

• existsok (bool, optional) – whether it is ok to overwrite an existing file. Defaults to False.

property filename

The filename of the tensor, if it has one.

Raises an exception otherwise.

classmethod from_filename(filename, dtype, shape, index=None)

Loads a MemoryMappedTensor from a given filename.

Parameters:

• filename (path or equivalent) – the path to the file.

• dtype (torch.dtype) – the dtype of the tensor.

• shape (torch.Size or torch.Tensor) – the shape of the tensor. If a tensor is provided, it is assumed that the tensor is a nested_tensor instance.

• index (torch-compatible index type) – an index to use to build the tensor.

classmethod from_handler(handler, dtype, shape, index=None)

Loads a MemoryMappedTensor from a given handler.

Parameters:

• handler (compatible file handler) – the handler for the tensor.

• dtype (torch.dtype) – the dtype of the tensor.

• shape (torch.Size or torch.Tensor) – the shape of the tensor. If a tensor is provided, it is assumed that the tensor is a nested_tensor instance.

• index (torch-compatible index type, optional) – an index to use to build the tensor.

classmethod from_tensor(input, *, filename: Optional[Union[Path, str]] = None, existsok: bool = False, copy_existing: bool = False, copy_data: bool = True, shape: Optional[Size] = None)

Creates a MemoryMappedTensor with the same content as another tensor.

If the tensor is already a MemoryMappedTensor the original tensor is returned if the filename argument is None or if the two paths match. In all other cases, a new MemoryMappedTensor is produced.

Parameters:

input (torch.Tensor) – the tensor which content must be copied onto the MemoryMappedTensor.

Keyword Arguments:

• filename (path to a file) – the path to the file where the tensor should be stored. If none is provided, a file handler is used instead.

• existsok (bool, optional) – if True, the file will overwrite an existing file. Defaults to False.

• copy_existing (bool, optional) – if True and the provided input is a MemoryMappedTensor with an associated filename, copying the content to the new location is permitted. Otherwise, an exception is thrown. This behaviour exists to prevent inadvertently duplicating data on disk.

• copy_data (bool, optional) – if True, the content of the tensor will be copied on the storage. Defaults to True.

• shape (torch.Size or torch.Tensor) – a shape to override the tensor shape. If a tensor is passed, it must represent the nested shapes of a nested tensor.

classmethod full(*size, fill_value, dtype=None, device=None, filename=None)

classmethod full(shape, *, fill_value, dtype=None, device=None, filename=None)

Creates a tensor with a single content specified by fill_value, specific shape, dtype and filename.

Parameters:

shape (integers or torch.Size) – the shape of the tensor.

Keyword Arguments:

• fill_value (float or equivalent) – content of the tensor.

• dtype (torch.dtype) – the dtype of the tensor.

• device (torch.device) – the device of the tensor. Only None and “cpu” are accepted, any other device will raise an exception.

• filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

• existsok (bool, optional) – whether it is ok to overwrite an existing file. Defaults to False.

classmethod full_like(input, fill_value, *, filename=None)

Creates a tensor with a single content indicated by the fill_value argument, but the same shape and dtype as the input tensor.

Parameters:

• input (torch.Tensor) – the tensor to use as an example.

• fill_value (float or equivalent) – content of the tensor.

Keyword Arguments:

filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

classmethod ones(*size, dtype=None, device=None, filename=None)

classmethod ones(shape, *, dtype=None, device=None, filename=None)

Creates a tensor with a 1-filled content, specific shape, dtype and filename.

Parameters:

shape (integers or torch.Size) – the shape of the tensor.

Keyword Arguments:

• dtype (torch.dtype) – the dtype of the tensor.

• device (torch.device) – the device of the tensor. Only None and “cpu” are accepted, any other device will raise an exception.

• filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

• existsok (bool, optional) – whether it is ok to overwrite an existing file. Defaults to False.

classmethod ones_like(input, *, filename=None)

Creates a tensor with a 1-filled content, but the same shape and dtype as the input tensor.

Parameters:

input (torch.Tensor) – the tensor to use as an example.

Keyword Arguments:

filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

unbind(dim=0) → seq

See torch.unbind()

classmethod zeros(*size, dtype=None, device=None, filename=None)

classmethod zeros(shape, *, dtype=None, device=None, filename=None)

Creates a tensor with a 0-filled content, specific shape, dtype and filename.

Parameters:

shape (integers or torch.Size) – the shape of the tensor.

Keyword Arguments:

• dtype (torch.dtype) – the dtype of the tensor.

• device (torch.device) – the device of the tensor. Only None and “cpu” are accepted, any other device will raise an exception.

• filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.

• existsok (bool, optional) – whether it is ok to overwrite an existing file. Defaults to False.

classmethod zeros_like(input, *, filename=None)

Creates a tensor with a 0-filled content, but the same shape and dtype as the input tensor.

Parameters:

input (torch.Tensor) – the tensor to use as an example.

Keyword Arguments:

filename (path or equivalent) – the path to the file, if any. If none is provided, a handler is used.