tensor.py 16.7 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
from textwrap import indent

import torch
import scipy.sparse

rusty1s's avatar
rusty1s committed
6
from torch_sparse.storage import SparseStorage, get_layout
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
from torch_sparse.transpose import t
rusty1s's avatar
rusty1s committed
9
from torch_sparse.narrow import narrow
rusty1s's avatar
rusty1s committed
10
11
12
from torch_sparse.select import select
from torch_sparse.index_select import index_select, index_select_nnz
from torch_sparse.masked_select import masked_select, masked_select_nnz
rusty1s's avatar
rusty1s committed
13
14
15
16


class SparseTensor(object):
    def __init__(self, index, value=None, sparse_size=None, is_sorted=False):
rusty1s's avatar
rusty1s committed
17
18
        self.storage = SparseStorage(index, value, sparse_size,
                                     is_sorted=is_sorted)
rusty1s's avatar
rusty1s committed
19
20
21
22

    @classmethod
    def from_storage(self, storage):
        self = SparseTensor.__new__(SparseTensor)
rusty1s's avatar
rusty1s committed
23
        self.storage = storage
rusty1s's avatar
rusty1s committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
        return self

    @classmethod
    def from_dense(self, mat):
        if mat.dim() > 2:
            index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero()
        else:
            index = mat.nonzero()

        index = index.t().contiguous()
        value = mat[index[0], index[1]]
        return self.__class__(index, value, mat.size()[:2], is_sorted=True)

    def __copy__(self):
rusty1s's avatar
rusty1s committed
38
        return self.from_storage(self.storage)
rusty1s's avatar
rusty1s committed
39
40

    def clone(self):
rusty1s's avatar
rusty1s committed
41
        return self.from_storage(self.storage.clone())
rusty1s's avatar
rusty1s committed
42
43
44
45
46
47
48
49
50

    def __deepcopy__(self, memo):
        new_sparse_tensor = self.clone()
        memo[id(self)] = new_sparse_tensor
        return new_sparse_tensor

    # Formats #################################################################

    def coo(self):
rusty1s's avatar
rusty1s committed
51
        return self.storage.index, self.storage.value
rusty1s's avatar
rusty1s committed
52
53

    def csr(self):
rusty1s's avatar
rusty1s committed
54
        return self.storage.rowptr, self.storage.col, self.storage.value
rusty1s's avatar
rusty1s committed
55
56

    def csc(self):
rusty1s's avatar
rusty1s committed
57
58
59
        perm = self.storage.csr2csc
        return (self.storage.colptr, self.storage.row[perm],
                self.storage.value[perm] if self.has_value() else None)
rusty1s's avatar
rusty1s committed
60
61
62
63

    # Storage inheritance #####################################################

    def has_value(self):
rusty1s's avatar
rusty1s committed
64
        return self.storage.has_value()
rusty1s's avatar
rusty1s committed
65
66

    def set_value_(self, value, layout=None):
rusty1s's avatar
rusty1s committed
67
        self.storage.set_value_(value, layout)
rusty1s's avatar
rusty1s committed
68
69
70
        return self

    def set_value(self, value, layout=None):
rusty1s's avatar
rusty1s committed
71
        return self.from_storage(self.storage.set_value(value, layout))
rusty1s's avatar
rusty1s committed
72
73

    def sparse_size(self, dim=None):
rusty1s's avatar
rusty1s committed
74
        return self.storage.sparse_size(dim)
rusty1s's avatar
rusty1s committed
75
76

    def sparse_resize_(self, *sizes):
rusty1s's avatar
rusty1s committed
77
        self.storage.sparse_resize_(*sizes)
rusty1s's avatar
rusty1s committed
78
79
80
        return self

    def is_coalesced(self):
rusty1s's avatar
rusty1s committed
81
        return self.storage.is_coalesced()
rusty1s's avatar
rusty1s committed
82
83

    def coalesce(self):
rusty1s's avatar
rusty1s committed
84
        return self.from_storage(self.storage.coalesce())
rusty1s's avatar
rusty1s committed
85
86

    def cached_keys(self):
rusty1s's avatar
rusty1s committed
87
        return self.storage.cached_keys()
rusty1s's avatar
rusty1s committed
88
89

    def fill_cache_(self, *args):
rusty1s's avatar
rusty1s committed
90
        self.storage.fill_cache_(*args)
rusty1s's avatar
rusty1s committed
91
92
93
        return self

    def clear_cache_(self, *args):
rusty1s's avatar
rusty1s committed
94
        self.storage.clear_cache_(*args)
rusty1s's avatar
rusty1s committed
95
96
97
98
99
100
        return self

    # Utility functions #######################################################

    def size(self, dim=None):
        size = self.sparse_size()
rusty1s's avatar
rusty1s committed
101
        size += self.storage.value.size()[1:] if self.has_value() else ()
rusty1s's avatar
rusty1s committed
102
103
104
105
106
107
108
109
110
111
        return size if dim is None else size[dim]

    def dim(self):
        return len(self.size())

    @property
    def shape(self):
        return self.size()

    def nnz(self):
rusty1s's avatar
rusty1s committed
112
        return self.storage.index.size(1)
rusty1s's avatar
rusty1s committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

    def density(self):
        return self.nnz() / (self.sparse_size(0) * self.sparse_size(1))

    def sparsity(self):
        return 1 - self.density()

    def avg_row_length(self):
        return self.nnz() / self.sparse_size(0)

    def avg_col_length(self):
        return self.nnz() / self.sparse_size(1)

    def numel(self):
        return self.value.numel() if self.has_value() else self.nnz()

    def is_quadratic(self):
        return self.sparse_size(0) == self.sparse_size(1)

    def is_symmetric(self):
        if not self.is_quadratic:
            return False

        rowptr, col, val1 = self.csr()
        colptr, row, val2 = self.csc()
rusty1s's avatar
rusty1s committed
138
139
140
        index_sym = (rowptr == colptr).all() and (col == row).all()
        value_sym = (val1 == val2).all().item() if self.has_value() else True
        return index_sym.item() and value_sym
rusty1s's avatar
rusty1s committed
141
142

    def detach_(self):
rusty1s's avatar
rusty1s committed
143
        self.storage.apply_(lambda x: x.detach_())
rusty1s's avatar
rusty1s committed
144
145
146
        return self

    def detach(self):
rusty1s's avatar
rusty1s committed
147
        return self.from_storage(self.storage.apply(lambda x: x.detach()))
rusty1s's avatar
rusty1s committed
148
149

    def pin_memory(self):
rusty1s's avatar
rusty1s committed
150
        return self.from_storage(self.storage.apply(lambda x: x.pin_memory()))
rusty1s's avatar
rusty1s committed
151
152

    def is_pinned(self):
rusty1s's avatar
rusty1s committed
153
        return all(self.storage.map(lambda x: x.is_pinned()))
rusty1s's avatar
rusty1s committed
154
155

    def share_memory_(self):
rusty1s's avatar
rusty1s committed
156
        self.storage.apply_(lambda x: x.share_memory_())
rusty1s's avatar
rusty1s committed
157
158
159
        return self

    def is_shared(self):
rusty1s's avatar
rusty1s committed
160
        return all(self.storage.map(lambda x: x.is_shared()))
rusty1s's avatar
rusty1s committed
161
162
163

    @property
    def device(self):
rusty1s's avatar
rusty1s committed
164
        return self.storage.index.device
rusty1s's avatar
rusty1s committed
165
166

    def cpu(self):
rusty1s's avatar
rusty1s committed
167
        return self.from_storage(self.storage.apply(lambda x: x.cpu()))
rusty1s's avatar
rusty1s committed
168
169

    def cuda(self, device=None, non_blocking=False, **kwargs):
rusty1s's avatar
rusty1s committed
170
171
172
        storage = self.storage.apply(
            lambda x: x.cuda(device, non_blocking, **kwargs))
        return self.from_storage(storage)
rusty1s's avatar
rusty1s committed
173
174
175

    @property
    def is_cuda(self):
rusty1s's avatar
rusty1s committed
176
        return self.storage.index.is_cuda
rusty1s's avatar
rusty1s committed
177
178
179

    @property
    def dtype(self):
rusty1s's avatar
rusty1s committed
180
        return self.storage.value.dtype if self.has_value() else None
rusty1s's avatar
rusty1s committed
181
182

    def is_floating_point(self):
rusty1s's avatar
rusty1s committed
183
        value = self.storage.value
rusty1s's avatar
rusty1s committed
184
185
186
187
188
189
190
191
192
        return self.has_value() and torch.is_floating_point(value)

    def type(self, dtype=None, non_blocking=False, **kwargs):
        if dtype is None:
            return self.dtype

        if dtype == self.dtype:
            return self

rusty1s's avatar
rusty1s committed
193
194
195
196
        storage = self.storage.apply_value(
            lambda x: x.type(dtype, non_blocking, **kwargs))

        return self.from_storage(storage)
rusty1s's avatar
rusty1s committed
197
198
199
200
201
202
203

    def to(self, *args, **kwargs):
        storage = None

        if 'device' in kwargs:
            device = kwargs['device']
            del kwargs['device']
rusty1s's avatar
rusty1s committed
204
            storage = self.storage.apply(lambda x: x.to(
rusty1s's avatar
rusty1s committed
205
206
207
208
                device, non_blocking=getattr(kwargs, 'non_blocking', False)))

        for arg in args[:]:
            if isinstance(arg, str) or isinstance(arg, torch.device):
rusty1s's avatar
rusty1s committed
209
                storage = self.storage.apply(lambda x: x.to(
rusty1s's avatar
rusty1s committed
210
211
212
213
                    arg, non_blocking=getattr(kwargs, 'non_blocking', False)))
                args.remove(arg)

        if storage is not None:
rusty1s's avatar
rusty1s committed
214
            self = self.from_storage(storage)
rusty1s's avatar
rusty1s committed
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

        if len(args) > 0 or len(kwargs) > 0:
            self = self.type(*args, **kwargs)

        return self

    def bfloat16(self):
        return self.type(torch.bfloat16)

    def bool(self):
        return self.type(torch.bool)

    def byte(self):
        return self.type(torch.byte)

    def char(self):
        return self.type(torch.char)

    def half(self):
        return self.type(torch.half)

    def float(self):
        return self.type(torch.float)

    def double(self):
        return self.type(torch.double)

    def short(self):
        return self.type(torch.short)

    def int(self):
        return self.type(torch.int)

    def long(self):
        return self.type(torch.long)

    # Conversions #############################################################

    def to_dense(self, dtype=None):
        dtype = dtype or self.dtype
        (row, col), value = self.coo()
        mat = torch.zeros(self.size(), dtype=dtype, device=self.device)
        mat[row, col] = value if self.has_value() else 1
        return mat

    def to_torch_sparse_coo_tensor(self, dtype=None, requires_grad=False):
        index, value = self.coo()
        return torch.sparse_coo_tensor(
rusty1s's avatar
rusty1s committed
263
264
265
266
267
            index, value if self.has_value() else torch.ones(
                self.nnz(), dtype=dtype, device=self.device), self.size(),
            device=self.device, requires_grad=requires_grad)

    def to_scipy(self, dtype=None, layout=None):
rusty1s's avatar
rusty1s committed
268
        assert self.dim() == 2
rusty1s's avatar
rusty1s committed
269
        layout = get_layout(layout)
rusty1s's avatar
rusty1s committed
270

rusty1s's avatar
rusty1s committed
271
272
        if not self.has_value():
            ones = torch.ones(self.nnz(), dtype=dtype).numpy()
rusty1s's avatar
rusty1s committed
273
274

        if layout == 'coo':
rusty1s's avatar
rusty1s committed
275
276
277
278
            (row, col), value = self.coo()
            row = row.detach().cpu().numpy()
            col = col.detach().cpu().numpy()
            value = value.detach().cpu().numpy() if self.has_value() else ones
rusty1s's avatar
rusty1s committed
279
280
            return scipy.sparse.coo_matrix((value, (row, col)), self.size())
        elif layout == 'csr':
rusty1s's avatar
rusty1s committed
281
282
283
284
            rowptr, col, value = self.csr()
            rowptr = rowptr.detach().cpu().numpy()
            col = col.detach().cpu().numpy()
            value = value.detach().cpu().numpy() if self.has_value() else ones
rusty1s's avatar
rusty1s committed
285
286
            return scipy.sparse.csr_matrix((value, col, rowptr), self.size())
        elif layout == 'csc':
rusty1s's avatar
rusty1s committed
287
288
289
290
            colptr, row, value = self.csc()
            colptr = colptr.detach().cpu().numpy()
            row = row.detach().cpu().numpy()
            value = value.detach().cpu().numpy() if self.has_value() else ones
rusty1s's avatar
rusty1s committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
            return scipy.sparse.csc_matrix((value, row, colptr), self.size())

    # String Reputation #######################################################

    def __repr__(self):
        i = ' ' * 6
        index, value = self.coo()
        infos = [f'index={indent(index.__repr__(), i)[len(i):]}']

        if self.has_value():
            infos += [f'value={indent(value.__repr__(), i)[len(i):]}']

        infos += [
            f'size={tuple(self.size())}, '
            f'nnz={self.nnz()}, '
            f'density={100 * self.density():.02f}%'
        ]
        infos = ',\n'.join(infos)

        i = ' ' * (len(self.__class__.__name__) + 1)
        return f'{self.__class__.__name__}({indent(infos, i)[len(i):]})'


# Bindings ####################################################################
rusty1s's avatar
rusty1s committed
315

rusty1s's avatar
rusty1s committed
316
SparseTensor.t = t
rusty1s's avatar
rusty1s committed
317
SparseTensor.narrow = narrow
rusty1s's avatar
rusty1s committed
318
319
320
321
322
SparseTensor.select = select
SparseTensor.index_select = index_select
SparseTensor.index_select_nnz = index_select_nnz
SparseTensor.masked_select = masked_select
SparseTensor.masked_select_nnz = masked_select_nnz
rusty1s's avatar
rusty1s committed
323
324
325
326
327
328

#     def __getitem__(self, idx):
#         # Convert int and slice to index tensor
#         # Filter list into edge and sparse slice
#         raise NotImplementedError

rusty1s's avatar
rusty1s committed
329
330
331
#     def set_diag(self, value):
#         raise NotImplementedError

rusty1s's avatar
rusty1s committed
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
#     def __reduce(self, dim, reduce, only_nnz):
#         raise NotImplementedError

#     def sum(self, dim):
#         return self.__reduce(dim, reduce='add', only_nnz=True)

#     def prod(self, dim):
#         return self.__reduce(dim, reduce='mul', only_nnz=True)

#     def min(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='min', only_nnz=only_nnz)

#     def max(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='min', only_nnz=only_nnz)

#     def mean(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='mean', only_nnz=only_nnz)

#     def matmul(self, mat, reduce='add'):
#         assert self.numel() == self.nnz()  # Disallow multi-dimensional value
#         if torch.is_tensor(mat):
#             raise NotImplementedError
#         elif isinstance(mat, self.__class__):
#             assert reduce == 'add'
rusty1s's avatar
rusty1s committed
356
#           assert mat.numel() == mat.nnz()  # Disallow multi-dimensional value
rusty1s's avatar
rusty1s committed
357
358
359
360
361
362
363
364
365
366
367
368
369
370
#             raise NotImplementedError
#         raise ValueError('Argument needs to be of type `torch.tensor` or '
#                          'type `torch_sparse.SparseTensor`.')

#     def add(self, other, layout=None):
#         if __is_scalar__(other):
#             if self.has_value:
#                 return self.set_value(self._value + other, 'coo')
#             else:
#                 return self.set_value(torch.full((self.nnz(), ), other + 1),
#                                       'coo')
#         elif torch.is_tensor(other):
#             if layout is None:
#                 layout = 'coo'
rusty1s's avatar
rusty1s committed
371
372
#               warnings.warn('`layout` argument unset, using default layout '
#                             '"coo". This may lead to unexpected behaviour.')
rusty1s's avatar
rusty1s committed
373
374
#             assert layout in ['coo', 'csr', 'csc']
#             if layout == 'csc':
rusty1s's avatar
rusty1s committed
375
#                 other = other[self._arg_csc2csr]
rusty1s's avatar
rusty1s committed
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
#             if self.has_value:
#                 return self.set_value(self._value + other, 'coo')
#             else:
#                 return self.set_value(other + 1, 'coo')
#         elif isinstance(other, self.__class__):
#             raise NotImplementedError
#         raise ValueError('Argument needs to be of type `int`, `float`, '
#                          '`torch.tensor` or `torch_sparse.SparseTensor`.')

#     def add_(self, other, layout=None):
#         if isinstance(other, int) or isinstance(other, float):
#             raise NotImplementedError
#         elif torch.is_tensor(other):
#             raise NotImplementedError
#         raise ValueError('Argument needs to be a scalar or of type '
#                          '`torch.tensor`.')

#     def __add__(self, other):
#         return self.add(other)

#     def __radd__(self, other):
#         return self.add(other)

#     def sub(self, layout=None):
#         raise NotImplementedError

#     def sub_(self, layout=None):
#         raise NotImplementedError

#     def mul(self, layout=None):
#         raise NotImplementedError

#     def mul_(self, layout=None):
#         raise NotImplementedError

#     def div(self, layout=None):
#         raise NotImplementedError

#     def div_(self, layout=None):
#         raise NotImplementedError

rusty1s's avatar
rusty1s committed
417
418
419
if __name__ == '__main__':
    from torch_geometric.datasets import Reddit, Planetoid  # noqa
    import time  # noqa
rusty1s's avatar
rusty1s committed
420

rusty1s's avatar
rusty1s committed
421
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
rusty1s's avatar
rusty1s committed
422

rusty1s's avatar
rusty1s committed
423
424
425
    # dataset = Reddit('/tmp/Reddit')
    dataset = Planetoid('/tmp/Cora', 'Cora')
    data = dataset[0].to(device)
rusty1s's avatar
rusty1s committed
426

rusty1s's avatar
rusty1s committed
427
    value = torch.randn((data.num_edges, 10), device=device)
rusty1s's avatar
rusty1s committed
428

rusty1s's avatar
rusty1s committed
429
    mat1 = SparseTensor(data.edge_index, value)
rusty1s's avatar
rusty1s committed
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467

    index = torch.tensor([0, 2])
    mat2 = mat1.index_select(2, index)

    index = torch.randperm(data.num_nodes)[:data.num_nodes - 500]
    mask = torch.zeros(data.num_nodes, dtype=torch.bool)
    mask[index] = True

    t = time.perf_counter()
    for _ in range(1000):
        mat2 = mat1.index_select(0, index)
    print(time.perf_counter() - t)

    t = time.perf_counter()
    for _ in range(1000):
        mat2 = mat1.masked_select(0, mask)
    print(time.perf_counter() - t)

    # mat2 = mat1.narrow(1, start=0, length=3)
    # print(mat2)

    # index = torch.randperm(data.num_nodes)
    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.index_select(0, index)
    # print(time.perf_counter() - t)

    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.index_select(1, index)
    # print(time.perf_counter() - t)
    # raise NotImplementedError

    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.t().index_select(0, index).t()
    # print(time.perf_counter() - t)

rusty1s's avatar
rusty1s committed
468
    # print(mat1)
rusty1s's avatar
rusty1s committed
469
470
471
472
473
474
475
476
477
478
479
480
481
    # mat1.index_select((0, 1), torch.tensor([0, 1, 2, 3, 4]))

    # print(mat3)
    # print(mat3.storage.rowcount)

    # print(mat1)

    # (row, col), value = mat1.coo()
    # mask = row < 3
    # t = time.perf_counter()
    # for _ in range(10000):
    #     mat2 = mat1.narrow(1, start=10, length=2690)
    # print(time.perf_counter() - t)
rusty1s's avatar
rusty1s committed
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498

    # # print(mat1.to_dense().size())
    # print(mat1.to_torch_sparse_coo_tensor().to_dense().size())
    # print(mat1.to_scipy(layout='coo').todense().shape)
    # print(mat1.to_scipy(layout='csr').todense().shape)
    # print(mat1.to_scipy(layout='csc').todense().shape)

    # print(mat1.is_quadratic())
    # print(mat1.is_symmetric())

    # print(mat1.cached_keys())
    # mat1 = mat1.t()
    # print(mat1.cached_keys())
    # mat1 = mat1.t()
    # print(mat1.cached_keys())
    # print('-------- NARROW ----------')

rusty1s's avatar
rusty1s committed
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
    # t = time.perf_counter()
    # for _ in range(100):
    #     out = mat1.narrow(dim=0, start=10, length=10)
    #     # out.storage.colptr
    # print(time.perf_counter() - t)
    # print(out)
    # print(out.cached_keys())

    # t = time.perf_counter()
    # for _ in range(100):
    #     out = mat1.narrow(dim=1, start=10, length=2000)
    #     # out.storage.colptr
    # print(time.perf_counter() - t)
    # print(out)
    # print(out.cached_keys())
rusty1s's avatar
rusty1s committed
514
515

    # mat1 = mat1.narrow(0, start=10, length=10)
rusty1s's avatar
rusty1s committed
516
    # mat1.storage._value = torch.randn(mat1.nnz(), 20)
rusty1s's avatar
rusty1s committed
517
518
519
    # print(mat1.coo()[1].size())
    # mat1 = mat1.narrow(2, start=10, length=10)
    # print(mat1.coo()[1].size())
rusty1s's avatar
rusty1s committed
520
521
#     mat1 = mat1.t()

rusty1s's avatar
rusty1s committed
522
#   mat2 = torch.sparse_coo_tensor(data.edge_index, torch.ones(data.num_edges),
rusty1s's avatar
rusty1s committed
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
#                                    device=device)
#     mat2 = mat2.coalesce()
#     mat2 = mat2.t().coalesce()

#     index1, value1 = mat1.coo()
#     index2, value2 = mat2._indices(), mat2._values()
#     assert torch.allclose(index1, index2)

#     out1 = mat1.to_dense()
#     out2 = mat2.to_dense()
#     assert torch.allclose(out1, out2)

#     out = 2 + mat1
#     print(out)

#     # mat1[1]
#     # mat1[1, 1]
#     # mat1[..., -1]
#     # mat1[:, -1]
#     # mat1[1:4, 1:4]
#     # mat1[torch.tensor([0, 1, 2])]