Bases: _ManagedTensor
Managed Tensor class
This class is a subclass of torch.Tensor
It overrides the torch_function method
to send the tensor to the device manager
Properties
pinned: bool - True if tensor is pinned to the device
Methods
pin
Pin the tensor to the device
unpin
Unpin the tensor from the device
cpu
Move the tensor to cpu
cuda
Move the tensor to cuda. If no args are passed, uses the device manager
to find a suitable device. If args are passed, uses the first arg as
the device to move the tensor to.
Device Manager specific args can be directly passed to the function
to
Move the tensor to the device. If string input cuda is used then use the
device manager to find a suitable device.
Device Manager specific args can be directly passed to the function.
Example:
import torch
from managed import ManagedTensor as mt
# Create a cuda tensor
a = torch.rand(10, 10).as_subclass(mt).cuda()
# Pin the tensor to gpu device
a.pin()
# Move the tensor to cpu
# Now pinned to cpu!
a = a.cpu()
# Create another cuda tensor
b = torch.rand(10, 10).as_subclass(mt).cuda()
# Result is on cpu as a is pinned to cpu
c = a + b
ManagedTensor can be mixed with normal torch.Tensor
Though it is recommended to use ManagedTensor for all tensors
as normal tensor will be changed to _ManagedTensor in future.
Example:
a = torch.rand(10, 10).as_subclass(mt).cuda()
b = torch.rand(10, 10)
# Class of b will be changed to _ManagedTensor
# Result and b will be on cuda device of a
c = a + b
Source code in managed/tensor.py
| class ManagedTensor(_ManagedTensor):
""" Managed Tensor class
This class is a subclass of torch.Tensor
It overrides the __torch_function__ method
to send the tensor to the device manager
## Properties
pinned: bool - True if tensor is pinned to the device
## Methods
pin:
Pin the tensor to the device
unpin:
Unpin the tensor from the device
cpu:
Move the tensor to cpu
cuda:
Move the tensor to cuda. If no args are passed, uses the device manager
to find a suitable device. If args are passed, uses the first arg as
the device to move the tensor to.
Device Manager specific args can be directly passed to the function
to:
Move the tensor to the device. If string input cuda is used then use the
device manager to find a suitable device.
Device Manager specific args can be directly passed to the function.
Example:
```python
import torch
from managed import ManagedTensor as mt
# Create a cuda tensor
a = torch.rand(10, 10).as_subclass(mt).cuda()
# Pin the tensor to gpu device
a.pin()
# Move the tensor to cpu
# Now pinned to cpu!
a = a.cpu()
# Create another cuda tensor
b = torch.rand(10, 10).as_subclass(mt).cuda()
# Result is on cpu as a is pinned to cpu
c = a + b
```
ManagedTensor can be mixed with normal torch.Tensor
Though it is recommended to use ManagedTensor for all tensors
as normal tensor will be changed to _ManagedTensor in future.
Example:
```python
a = torch.rand(10, 10).as_subclass(mt).cuda()
b = torch.rand(10, 10)
# Class of b will be changed to _ManagedTensor
# Result and b will be on cuda device of a
c = a + b
```
"""
@classmethod
def __torch_function__(cls, func, types, args=[], kwargs=None):
if kwargs is None:
kwargs = {}
# TODO: This needs to be optimized
tensor_list = []
if func.__name__ not in FUNC_BLACKLIST:
aggregate_tensors(tensor_list, args)
aggregate_tensors(tensor_list, kwargs)
device_manager.send(tensor_list)
############################################
# TODO: This is a temporary fix for
# device type check from pytroch but it is not
# fool-proof. Rest require manual pinning of module
# Issue: https://github.com/pytorch/pytorch/issues/65016
# Remove this when issue is fixed properly
############################################
ret = super().__torch_function__(func, types, args, kwargs)
if func.__name__ not in FUNC_BLACKLIST and func.__name__ != "backward":
ret_list = []
aggregate_tensors(ret_list, ret)
if len(ret_list) == 0:
return ret
for t in ret_list:
hook_unexplored_graph(t.grad_fn, t.device)
return ret
def cuda(self, *args, **kwargs):
if len(args) > 0:
if isinstance(args[0], int):
return super().cuda(*args, **kwargs).as_subclass(self.__class__)
elif args[0] is None:
device = device_manager.cuda(self, *args, **kwargs)
elif isinstance(args[0], torch.device):
device = args[0]
else:
raise TypeError("Given device is not a cuda index or torch device")
elif 'device' in kwargs:
device = kwargs['device']
else:
device = device_manager.cuda(self, *args, **kwargs)
return super().to(device).as_subclass(self.__class__)
|