1import functools 2import operator 3import random 4import time 5 6import numpy as np 7 8import torch 9 10 11# shim for torch.cuda.Event when running on cpu 12class Event: 13 def __init__(self, enable_timing): 14 pass 15 16 def record(self): 17 self.time = time.perf_counter() 18 19 def elapsed_time(self, end_event): 20 assert isinstance(end_event, Event) 21 return end_event.time - self.time 22 23 24def gen_sparse_csr(shape, nnz): 25 fill_value = 0 26 total_values = functools.reduce(operator.mul, shape, 1) 27 dense = np.random.randn(total_values) 28 fills = random.sample(list(range(total_values)), total_values - nnz) 29 30 for f in fills: 31 dense[f] = fill_value 32 dense = torch.from_numpy(dense.reshape(shape)) 33 34 return dense.to_sparse_csr() 35 36 37def gen_sparse_coo(shape, nnz): 38 dense = np.random.randn(*shape) 39 values = [] 40 indices = [[], []] 41 for n in range(nnz): 42 row = random.randint(0, shape[0] - 1) 43 col = random.randint(0, shape[1] - 1) 44 indices[0].append(row) 45 indices[1].append(col) 46 values.append(dense[row, col]) 47 48 return torch.sparse_coo_tensor(indices, values, size=shape) 49 50 51def gen_sparse_coo_and_csr(shape, nnz): 52 total_values = functools.reduce(operator.mul, shape, 1) 53 dense = np.random.randn(total_values) 54 fills = random.sample(list(range(total_values)), total_values - nnz) 55 56 for f in fills: 57 dense[f] = 0 58 59 dense = torch.from_numpy(dense.reshape(shape)) 60 return dense.to_sparse(), dense.to_sparse_csr() 61