wehub-resource-sync
70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Implementation of a compressed sparse tensor. Similar in
|
|
functionality to TensorFlow's IndexedSlices implementation.
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
class SparseTensor(object):
|
|
""" Compressed Sparse Tensor """
|
|
|
|
def __init__(self, dense_tensor=None):
|
|
self.orig_dense_tensor = dense_tensor
|
|
if dense_tensor is not None:
|
|
self.is_sparse = dense_tensor.is_sparse
|
|
self.dtype = self.orig_dense_tensor.dtype
|
|
if self.is_sparse:
|
|
dense_tensor = dense_tensor.coalesce()
|
|
self.indices = dense_tensor.indices().flatten()
|
|
self.values = dense_tensor.values()
|
|
else:
|
|
result = torch.sum(dense_tensor, dim=1)
|
|
self.indices = result.nonzero().flatten()
|
|
self.values = dense_tensor[self.indices]
|
|
self.dense_size = list(dense_tensor.size())
|
|
else:
|
|
self.indices = None
|
|
self.values = None
|
|
self.dense_size = None
|
|
|
|
def to_coo_tensor(self):
|
|
return torch.sparse_coo_tensor(self.indices.unsqueeze(0), self.values, self.dense_size)
|
|
|
|
@staticmethod
|
|
def type():
|
|
return "deepspeed.SparseTensor"
|
|
|
|
def to_dense(self):
|
|
it = self.indices.unsqueeze(1)
|
|
full_indices = torch.cat([it for _ in range(self.dense_size[1])], dim=1)
|
|
return self.values.new_zeros(self.dense_size).scatter_add_(0, full_indices, self.values)
|
|
|
|
def sparse_size(self):
|
|
index_size = list(self.indices.size())
|
|
index_size = index_size[0]
|
|
value_size = list(self.values.size())
|
|
value_size = value_size[0] * value_size[1]
|
|
dense_size = self.dense_size[0] * self.dense_size[1]
|
|
return index_size + value_size, dense_size
|
|
|
|
def add(self, b):
|
|
assert self.dense_size == b.dense_size
|
|
self.indices = torch.cat([self.indices, b.indices])
|
|
self.values = torch.cat([self.values, b.values])
|
|
|
|
def __str__(self):
|
|
sparse_size, dense_size = self.sparse_size()
|
|
return "DeepSpeed.SparseTensor(indices_size={}, values_size={}, " \
|
|
"dense_size={}, device={}, reduction_factor={})".format(
|
|
self.indices.size(), self.values.size(), self.dense_size,
|
|
self.indices.get_device(), dense_size / sparse_size
|
|
)
|
|
|
|
def __repr__(self):
|
|
return self.__str__()
|