Source code for brainspy.utils.pytorch
"""
Class to consistently manage torch variables and their data types. It can format
data provided as a torch tensor,numpy array or an nn.Module object. It also han-
dles the intialization of seed for generation of random values.
"""
import torch
import numpy as np
import random
[docs]
class TorchUtils:
""" Consistently manage data movement and formatting of pytorch tensors and torch.nn.Module
instances. This includes data movement to/from CUDA and CPU, as well as transformations to/
from lists and numpy arrays. It also enables to manage the seeds for reproducibility purpo-
ses. More information can be found at:
https://pytorch.org/docs/stable/torch.html
https://pytorch.org/docs/stable/notes/randomness.html
"""
force_cpu = False
[docs]
@staticmethod
def set_force_cpu(force: bool):
"""
Facilitates running all tensors on the CPU even when having a GPU with cuda correctly
installed. Ideal for GPUs that do not have enough memory to run certain experiments or
for computers with or with older GPUs. Also, for cases where torch detects that there
is cuda, but the version is too old to be compatible.
Parameters
----------
force : boolean
True or false to set the force_cpu option to detect cuda.
Example
--------
TorchUtils.set_force_cpu(True)
"""
assert (type(force) == bool)
TorchUtils.force_cpu = force
[docs]
@staticmethod
def get_device():
"""
Consistently returns the device type for the torch. The device type can be "cpu" or "cuda",
depending on if the computer has a GPU with a cuda installation, and on the variable
force_cpu. This method does not support the management of multiple GPUs.
Returns
-------
torch.device
Device type of torch tensor which can be "cpu" or "cuda" depending on computer version.
Example
--------
TorchUtils.get_device()
"""
if torch.cuda.is_available() and not TorchUtils.force_cpu:
return torch.device("cuda")
return torch.device("cpu")
[docs]
@staticmethod
def format(data, device=None, data_type=None):
"""
If a list or a numpy array is provided, it enables to create a torch variable with a
consistent accelerator type and data type. If an exisiting torch tensor is provided,
the function enables setting the data type and device consistently for all torch tensors.
For devices with more than one GPU, if an instance of an nn.Module is provided, for example
CustomModel,DNPU or Processor, this function helps to distribute the model.
Parameters
----------
data : Data to be formatted. It can be one of these data types:
list : list of data indices
np.ndarray : list of data indices
torch.Tensor : inital torch tensor which has to be formatted
nn.Module : model of an nn.Module object
device : torch.device, optional
Device type of torch tensor which can be "cpu or "cuda" depending on computer version.
When set to none, it will take the default value from the global variable force_cpu of
this class. By default device is set to None.
data_type : torch.dtype, optional
desired data type of torch tensor, by default None
When set to None, it will take the default data type from pytorch. This datatype can be
changed using the torch.set_default_dtype. More info on this method at:
https://pytorch.org/docs/stable/generated/torch.set_default_dtype.html
Returns
-------
1. torch.tensor
Torch tensor either generated from python list or numpy array, or exisiting torch
tensors formatted to given device and data type.
2. nn.Module
If an nn.Module is given as an argument, a model of the nn.Module object distributed by
DataParallel amongst the multiple GPUs is generated.
Examples
-------
1. data = [[1, 2]]
tensor = TorchUtils.format(data, device=torch.device('cpu'))
2. tensor = torch.randn(2, 2)
tensor = TorchUtils.format(tensor, data_type=torch.float64)
3. data = [[1, 2], [3, 4]]
numpy_data = np.array(data)
tensor = TorchUtils.format(numpy_data)
4. model = CustomModel() # Where CustomModel is an instance of torch.nn.Module
newmodel = format_model(model)
"""
if device is None:
device = TorchUtils.get_device()
if data_type is None:
data_type = torch.get_default_dtype()
if isinstance(data, (np.ndarray, np.generic)):
return torch.tensor(data, device=device, dtype=data_type)
elif isinstance(data, torch.Tensor):
return data.to(device=device, dtype=data_type)
elif isinstance(data, list):
return torch.tensor(np.array(data), device=device, dtype=data_type)
elif isinstance(data, torch.nn.Module):
if torch.cuda.device_count() > 1 and not TorchUtils.force_cpu:
data = torch.nn.DataParallel(data)
data.to(TorchUtils.get_device())
return data
else:
raise TypeError(
"Input type not recognised. " +
"Supported types are: lists, Numpy arrays, Pytorch tensors and modules."
)
[docs]
@staticmethod
def to_numpy(data: torch.Tensor):
"""
Transforms torch tensor into a numpy array, detatching it first, and
moving the data to the cpu (if needed). The aim is to simplify the lines
of code required for this purpose with the original pytorch library.
Parameters
----------
data : torch.tensor
Torch tensor that needs to be formatted to a numpy array.
Returns
-------
np.array
Numpy array from given torch tensor.
Example
-------
tensor = torch.tensor([[1., -1.], [1., -1.]])
numpy_data = TorchUtils.to_numpy(tensor)
"""
if data.requires_grad:
return data.detach().cpu().numpy()
return data.cpu().numpy()
[docs]
@staticmethod
def init_seed(seed=None, deterministic=False):
"""
It enables to set a fixed seed, or retrieve the current seed used for generating random
numbers. If the random seed is not reset, different numbers appear with every invocation.
It can be used to reproduce results.
More information at: https://pytorch.org/docs/stable/notes/randomness.html
Parameters
----------
seed : int, optional
Value of seed, by default None.
deterministic : bool, optional
If the random value should be deterministic, by default False.
Returns
-------
int
Value of the seed that is being used.
Example
-------
TorchUtils.init_seed(0)
random1 = np.random.rand(4)
"""
if seed is None:
seed = random.randint(0, (2**32) - 1)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
if deterministic:
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
return seed