| |
| |
|
|
| |
| """ |
| 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 |
| self.dtype = self.orig_dense_tensor.dtype |
| self.is_sparse = dense_tensor.is_sparse |
| if dense_tensor is not None: |
| if dense_tensor.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__() |
|
|
