def _take_tensors(tensors, size_limit):
"""Group tensors into chunks. This generator yields a chunk at each time,
each containing tensors of same type up to certain byte limit in total size.
Args:
tensors (Sequence): A sequence of tensors to be separated into chunks.
size_limit (int): The limit of each chunk in bytes.
Yields:
Blocks of tensors of same type and within size_limit. The yielded
tensors are only ordered as the original sequence within its types.
"""
buf_dict = defaultdict(lambda: [[], 0])
for tensor in tensors:
t = tensor.type()
if tensor.is_sparse:
indices = torch._indices(tensor)
values = torch._values(tensor)
size = indices.numel() * indices.element_size() + values.numel() * values.element_size()
else:
size = tensor.numel() * tensor.element_size()
buf_and_size = buf_dict[t]
if buf_and_size[1] + size > size_limit and buf_and_size[1] > 0:
yield buf_and_size[0]
buf_and_size = buf_dict[t] = [[], 0]
buf_and_size[0].append(tensor)
buf_and_size[1] += size
for buf, _ in buf_dict.values():
if len(buf) > 0:
yield buf
def _cuda(self, device=None, non_blocking=False, **kwargs):
"""Returns a copy of this object in CUDA memory.
If this object is already in CUDA memory and on the correct device, then
no copy is performed and the original object is returned.
Args:
device (int): The destination GPU id. Defaults to the current device.
non_blocking (bool): If ``True`` and the source is in pinned memory,
the copy will be asynchronous with respect to the host. Otherwise,
the argument has no effect.
**kwargs: For compatibility, may contain the key ``async`` in place of
the ``non_blocking`` argument.
"""
non_blocking = _get_async_or_non_blocking('cuda', non_blocking, kwargs)
if self.is_cuda:
if device is None:
device = torch.cuda.current_device()
if self.get_device() == device:
return self
else:
if device is None:
device = -1
with torch.cuda.device(device):
if self.is_sparse:
new_type = getattr(torch.cuda.sparse, self.__class__.__name__)
indices = torch._indices(self).cuda(device, non_blocking)
values = torch._values(self).cuda(device, non_blocking)
return new_type(indices, values, self.size())
else:
new_type = getattr(torch.cuda, self.__class__.__name__)
return new_type(self.size()).copy_(self, non_blocking)
def _flatten_sparse_tensors(tensors):
"""Flatten sparse tensors into two contiguous 1D buffers, one of indices and
one of values. Assume tensors are of same sparse type.
Arguments:
tensors (Iterable[Tensor]): sparse tensors to flatten.
Returns:
A tuple of two contiguous 1D buffers, one containing input tensors'
indices and the other containing the values.
"""
flat_indices = _flatten_dense_tensors([torch._indices(t) for t in tensors])
flat_values = _flatten_dense_tensors([torch._values(t) for t in tensors])
return flat_indices, flat_values
def _unflatten_sparse_tensors(flat, tensors):
"""View flat buffer (containing indices and values) using the sizes of
tensors. Assume that tensors are of same sparse type, and that flat is given
by _flatten_sparse_tensors.
Arguments:
flat (tuple(Tensor, Tensor)): flattened indices and values of sparse
tensors to unflatten.
tensors (Iterable[Tensor]): sparse tensors whose sizes will be used to
unflatten flat.
Returns:
Unflattened sparse tensors with sizes same as tensors and values from
flat.
"""
flat_indices, flat_values = flat
indices = _unflatten_dense_tensors(flat_indices, [torch._indices(t) for t in tensors])
values = _unflatten_dense_tensors(flat_values, [torch._values(t) for t in tensors])
outputs = []
for t, i, v in zip(tensors, indices, values):
outputs.append(t.new(i, v, t.size()))
return tuple(outputs)
def _type(self, dtype=None, non_blocking=False, **kwargs):
"""Returns the type if `dtype` is not provided, else casts this object to
the specified type.
If this is already of the correct type, no copy is performed and the
original object is returned.
Args:
dtype (type or string): The desired type
non_blocking (bool): If ``True``, and the source is in pinned memory
and destination is on the GPU or vice versa, the copy is performed
asynchronously with respect to the host. Otherwise, the argument
has no effect.
**kwargs: For compatibility, may contain the key ``async`` in place of
the ``non_blocking`` argument. The ``async`` arg is deprecated.
"""
non_blocking = _get_async_or_non_blocking('type', non_blocking, kwargs)
if dtype is None:
return self.__module__ + '.' + self.__class__.__name__
if isinstance(dtype, str):
dtype = _import_dotted_name(dtype)
if dtype == type(self):
return self
if self.is_sparse:
if not dtype.is_sparse:
raise RuntimeError("Cannot cast sparse tensor to dense tensor")
new_module_name = dtype.__module__.replace('.sparse', '')
new_values_type_name = new_module_name + '.' + dtype.__name__
new_values = torch._values(self).type(new_values_type_name,
non_blocking)
new_indices_type_name = new_module_name + '.LongTensor'
new_indices = torch._indices(self).type(new_indices_type_name,
non_blocking)
return dtype(new_indices, new_values, self.size())
if dtype.is_sparse:
raise RuntimeError("Cannot cast dense tensor to sparse tensor")
return dtype(self.size()).copy_(self, non_blocking)