def bdsmm(sparse, dense):
"""
Batch dense-sparse matrix multiply
"""
if sparse.ndimension() > 2:
batch_size, n_rows, n_cols = sparse.size()
batch_assignment = sparse._indices()[0]
indices = sparse._indices()[1:].clone()
indices[0].add_(n_rows, batch_assignment)
indices[1].add_(n_cols, batch_assignment)
sparse_2d = sparse.__class__(indices, sparse._values(),
torch.Size((batch_size * n_rows, batch_size * n_cols)))
if dense.size(0) == 1:
dense = dense.repeat(batch_size, 1, 1)
dense_2d = dense.contiguous().view(batch_size * n_cols, -1)
res = torch.dsmm(sparse_2d, dense_2d)
res = res.view(batch_size, n_rows, -1)
return res
else:
return torch.dsmm(sparse, dense)
python类sparse()的实例源码
def cpu(self):
return self.type(getattr(torch.sparse, self.__class__.__name__))
def __init__(self, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
super(Embedding, self).__init__()
self.padding_idx = padding_idx
self.max_norm = max_norm
self.norm_type = norm_type
self.scale_grad_by_freq = scale_grad_by_freq
self._indices = None
self.sparse = sparse
def setUp(self):
# These parameters control the various ways we can run the test.
# We will subclass and override this method to implement CUDA
# tests
self.is_cuda = False
self.is_uncoalesced = False
self.IndexTensor = torch.LongTensor
self.ValueTensor = torch.DoubleTensor
self.SparseTensor = torch.sparse.DoubleTensor
def _gen_sparse(self, d, nnz, with_size):
# TODO: Consider implementing this in the CUDA case by directly
# performing the operations on the GPU. You won't be able to
# use torch.rand/torch.randn in this case because they are
# CPU-only. If you do this, you can remove the is_cuda branch
# at the end.
#
# If you do this, be sure to update assert_uncoalesced too
if isinstance(with_size, Number):
with_size = [with_size] * d
if self.is_uncoalesced:
# We want to generate a tensor with a lot of uncoalesced
# entries to stress test whether or not we handle this
# (subtle) case correctly
v_size = [nnz * 2] + list(with_size[d:])
v = torch.randn(*v_size)
r = torch.rand(d, nnz)
# Repeat the indexes, so every position shows up twice
i = torch.cat([r, r], dim=1) * \
torch.Tensor(with_size[:d]).repeat(nnz * 2, 1).transpose(0, 1)
i = i.type(torch.LongTensor)
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
self.assert_uncoalesced(x)
else:
# Generate a sparse tensor with d sparse dimensions; the
# rest the dimensions with_size[d:] are dense.
v_size = [nnz] + list(with_size[d:])
v = torch.randn(*v_size)
i = torch.rand(d, nnz) * \
torch.Tensor(with_size[:d]).repeat(nnz, 1).transpose(0, 1)
i = i.type(torch.LongTensor)
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
if self.is_cuda:
return x.cuda(), i.cuda(), v.cuda()
else:
return x, i.clone(), v.clone()
def setUp(self):
super(TestCudaSparse, self).setUp()
self.is_cuda = True
self.IndexTensor = torch.cuda.LongTensor
self.ValueTensor = torch.cuda.DoubleTensor
self.SparseTensor = torch.cuda.sparse.DoubleTensor
def cpu(self):
return self.type(getattr(torch.sparse, self.__class__.__name__))
def __init__(self, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
super(Embedding, self).__init__()
self.padding_idx = padding_idx
self.max_norm = max_norm
self.norm_type = norm_type
self.scale_grad_by_freq = scale_grad_by_freq
self._indices = None
self.sparse = sparse
def setUp(self):
# These parameters control the various ways we can run the test.
# We will subclass and override this method to implement CUDA
# tests
self.is_cuda = False
self.is_uncoalesced = False
self.IndexTensor = torch.LongTensor
self.ValueTensor = torch.DoubleTensor
self.SparseTensor = torch.sparse.DoubleTensor
def _gen_sparse(self, d, nnz, with_size):
# TODO: Consider implementing this in the CUDA case by directly
# performing the operations on the GPU. You won't be able to
# use torch.rand/torch.randn in this case because they are
# CPU-only. If you do this, you can remove the is_cuda branch
# at the end.
#
# If you do this, be sure to update assert_uncoalesced too
if isinstance(with_size, Number):
with_size = [with_size] * d
if self.is_uncoalesced:
# We want to generate a tensor with a lot of uncoalesced
# entries to stress test whether or not we handle this
# (subtle) case correctly
v_size = [nnz * 2] + list(with_size[d:])
v = torch.randn(*v_size)
r = torch.rand(d, nnz)
# Repeat the indexes, so every position shows up twice
i = torch.cat([r, r], dim=1) * \
torch.Tensor(with_size[:d]).repeat(nnz * 2, 1).transpose(0, 1)
i = i.type(torch.LongTensor)
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
self.assert_uncoalesced(x)
else:
# Generate a sparse tensor with d sparse dimensions; the
# rest the dimensions with_size[d:] are dense.
v_size = [nnz] + list(with_size[d:])
v = torch.randn(*v_size)
i = torch.rand(d, nnz) * \
torch.Tensor(with_size[:d]).repeat(nnz, 1).transpose(0, 1)
i = i.type(torch.LongTensor)
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
if self.is_cuda:
return x.cuda(), i.cuda(), v.cuda()
else:
return x, i.clone(), v.clone()
def setUp(self):
super(TestCudaSparse, self).setUp()
self.is_cuda = True
self.IndexTensor = torch.cuda.LongTensor
self.ValueTensor = torch.cuda.DoubleTensor
self.SparseTensor = torch.cuda.sparse.DoubleTensor
def cpu(self):
return self.type(getattr(torch.sparse, self.__class__.__name__))
def symbolic(g, indices, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
if max_norm is not None:
raise ValueError('Right now, re-norm is not supported.')
output = g.appendNode(g.create("Gather", [weight, indices]))
return output
def forward(cls, ctx, indices, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
ctx.padding_idx = padding_idx
ctx.scale_grad_by_freq = scale_grad_by_freq
ctx._indices = None
ctx.sparse = sparse
assert indices.dim() <= 2
assert not ctx.needs_input_grad[0], "Embedding doesn't " \
"compute the gradient w.r.t. the indices"
ctx._backend = type2backend[type(weight)]
ctx._weight_size = weight.size()
if not indices.is_contiguous():
ctx._indices = indices.contiguous()
indices = ctx._indices
else:
ctx.save_for_backward(indices)
output = weight.new()
if max_norm is not None:
cls._renorm(ctx, indices, weight, max_norm, norm_type)
if indices.dim() == 1:
output = torch.index_select(weight, 0, indices)
else:
output = torch.index_select(weight, 0, indices.view(-1))
output = output.view(indices.size(0), indices.size(1), weight.size(1))
return output
def setUp(self):
# These parameters control the various ways we can run the test.
# We will subclass and override this method to implement CUDA
# tests
self.is_cuda = False
self.is_uncoalesced = False
self.IndexTensor = torch.LongTensor
self.ValueTensor = torch.DoubleTensor
self.SparseTensor = torch.sparse.DoubleTensor
def assert_uncoalesced(self, x):
"""
Test if a CPU tensor is uncoalesced. This is used to ensure
correctness of the uncoalesced tensor generation algorithm.
"""
assert not x.is_coalesced()
# Strategy: construct a new sparse tensor with the raw value
# field overwritten to a tensor of ones, coalesce it, and then
# check if any value entries are > 1 (which indicates that the
# original was uncoalesced.)
i = x._indices().clone()
v = x._values().clone().fill_(1)
y = torch.sparse.DoubleTensor(i, v, x.size())
z = self.safeCoalesce(y)
assert (z._values() > 1).sum() > 0
def test_storage_not_null(self):
x = torch.cuda.sparse.FloatTensor(2)
self.assertNotEqual(x.get_device(), -1)
def setUp(self):
super(TestCudaSparse, self).setUp()
self.is_cuda = True
self.IndexTensor = torch.cuda.LongTensor
self.ValueTensor = torch.cuda.DoubleTensor
self.SparseTensor = torch.cuda.sparse.DoubleTensor
def sparse_eye(size):
"""
Returns the identity matrix as a sparse matrix
"""
indices = torch.arange(0, size).long().unsqueeze(0).expand(2, size)
values = torch.Tensor([1]).expand(size)
return torch.sparse.FloatTensor(indices, values, torch.Size([size, size]))
def sparse_repeat(sparse, *repeat_sizes):
orig_ndim = sparse.ndimension()
new_ndim = len(repeat_sizes)
orig_nvalues = sparse._indices().size(1)
# Expand the number of dimensions to match repeat_sizes
indices = torch.cat([sparse._indices().new().resize_(new_ndim - orig_ndim, orig_nvalues).zero_(),
sparse._indices()])
values = sparse._values()
size = [1] * (new_ndim - orig_ndim) + list(sparse.size())
# Expand each dimension
new_indices = indices.new().resize_(indices.size(0), indices.size(1) * mul(*repeat_sizes)).zero_()
new_values = values.repeat(mul(*repeat_sizes))
new_size = [dim_size * repeat_size for dim_size, repeat_size in zip(size, repeat_sizes)]
# Fill in new indices
new_indices[:, :orig_nvalues].copy_(indices)
unit_size = orig_nvalues
for i in range(new_ndim)[::-1]:
repeat_size = repeat_sizes[i]
for j in range(1, repeat_size):
new_indices[:, unit_size * j:unit_size * (j + 1)].copy_(new_indices[:, :unit_size])
new_indices[i, unit_size * j:unit_size * (j + 1)] += j * size[i]
unit_size *= repeat_size
return sparse.__class__(new_indices, new_values, torch.Size(new_size))
def to_sparse(dense):
mask = dense.ne(0)
indices = mask.nonzero()
if indices.storage():
values = dense[mask]
else:
indices = indices.resize_(1, dense.ndimension()).zero_()
values = dense.new().resize_(1).zero_()
# Construct sparse tensor
klass = getattr(torch.sparse, dense.__class__.__name__)
res = klass(indices.t(), values, dense.size())
if dense.is_cuda:
res = res.cuda()
return res
def cpu(self):
return self.type(getattr(torch.sparse, self.__class__.__name__))
def symbolic(g, indices, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
if max_norm is not None:
raise ValueError('Right now, re-norm is not supported.')
return g.op("Gather", weight, indices)
def forward(cls, ctx, indices, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq,
sparse=False):
ctx.padding_idx = padding_idx
ctx.scale_grad_by_freq = scale_grad_by_freq
ctx._indices = None
ctx.sparse = sparse
assert indices.dim() <= 2
assert not ctx.needs_input_grad[0], "Embedding doesn't " \
"compute the gradient w.r.t. the indices"
ctx._backend = type2backend[type(weight)]
ctx._weight_size = weight.size()
if not indices.is_contiguous():
ctx._indices = indices.contiguous()
indices = ctx._indices
else:
ctx.save_for_backward(indices)
output = weight.new()
if max_norm is not None:
cls._renorm(ctx, indices, weight, max_norm, norm_type)
if indices.dim() == 1:
output = torch.index_select(weight, 0, indices)
else:
output = torch.index_select(weight, 0, indices.view(-1))
output = output.view(indices.size(0), indices.size(1), weight.size(1))
return output
def setUp(self):
# These parameters control the various ways we can run the test.
# We will subclass and override this method to implement CUDA
# tests
self.is_cuda = False
self.is_uncoalesced = False
self.IndexTensor = torch.LongTensor
self.ValueTensor = torch.DoubleTensor
self.SparseTensor = torch.sparse.DoubleTensor
super(TestSparse, self).setUp()
def assert_uncoalesced(self, x):
"""
Test if a CPU tensor is uncoalesced. This is used to ensure
correctness of the uncoalesced tensor generation algorithm.
"""
assert not x.is_coalesced()
# Strategy: construct a new sparse tensor with the raw value
# field overwritten to a tensor of ones, coalesce it, and then
# check if any value entries are > 1 (which indicates that the
# original was uncoalesced.)
i = x._indices().clone()
v = x._values().clone().fill_(1)
y = torch.sparse.DoubleTensor(i, v, x.size())
z = self.safeCoalesce(y)
assert (z._values() > 1).sum() > 0
def test_storage_not_null(self):
x = torch.cuda.sparse.FloatTensor(2)
self.assertNotEqual(x.get_device(), -1)
def _test_new_device(self, size, device):
with torch.cuda.device(device):
x = torch.cuda.sparse.DoubleTensor(*size)
self.assertEqual(x.get_device(), device)
x1 = x.new()
x2 = x.new(2, 3)
self.assertEqual(x1.get_device(), device)
self.assertEqual(x2.get_device(), device)
def setUp(self):
super(TestCudaSparse, self).setUp()
self.is_cuda = True
self.IndexTensor = torch.cuda.LongTensor
self.ValueTensor = torch.cuda.DoubleTensor
self.SparseTensor = torch.cuda.sparse.DoubleTensor
def backward(self, grad_output):
if self._indices is not None:
indices = self._indices
else:
indices, = self.saved_tensors
grad_output = grad_output.contiguous()
if not self.sparse:
if indices.dim() == 2:
indices = indices.view(-1)
with torch.cuda.device_of(grad_output):
if grad_output.is_cuda:
_sorted = torch.cuda.LongTensor()
_indices = torch.cuda.LongTensor()
_count = torch.cuda.LongTensor()
else:
_count = torch.IntTensor()
_sorted = _indices = None
grad_weight = grad_output.new(self._weight_size).zero_()
self._backend.LookupTable_accGradParameters(
self._backend.library_state,
indices,
grad_output,
grad_weight,
_count,
_sorted,
_indices,
self.scale_grad_by_freq,
self.padding_idx,
1
)
else:
tensor_type = type(grad_output).__name__
if grad_output.is_cuda:
SparseTensor = getattr(torch.cuda.sparse, tensor_type)
else:
SparseTensor = getattr(torch.sparse, tensor_type)
grad_weight = SparseTensor(
indices.view(1, -1),
grad_output.view(-1, self._weight_size[1]),
self._weight_size,
)
return None, grad_weight
