Tools: Memory

Tools: Memory#

Functions for getting information on memory requirements for tensors and devices (CPU/GPU).

Example

>>> import torch
>>> from tad_mctc.tools import memory_tensor
>>> size = (100, 100)
>>> dtype = torch.float32
>>> memory_tensor(size, dtype)
0.03814697265625
tad_mctc.tools.memory.memory_device(device)[source]#

Get the available and total memory of the device.

Parameters:

device (torch.device) – Device to check memory for.

Returns:

Available and total memory in MB.

Return type:

tuple[float, float]

tad_mctc.tools.memory.memory_tensor(size, dtype)[source]#

Estimate the memory usage of a tensor.

Parameters:

  • size (Size) – Shape of the tensor.

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

Returns:

Estimated memory usage in MB.

Return type:

float

Raises:

ValueError – If the tensor data type is not supported.

Contents