/
vshmidt
/
pytorch
Обзор
Документация
Войти
/
vshmidt
/
pytorch
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
benchmarks/sparse/utils.py
60 строк
1 KB
Xuehai Pan
[5/N][Easy] fix typo for `usort` config in `pyproject.toml` (`kown` -> `known`): sort torch (#127126)
27 май 2024, 17:49
27 май 2024, 17:49
26f4f10
Код
Авторство
О чём код?
import functools import operator import random import time import numpy as np import torch # shim for torch.cuda.Event when running on cpu class Event: def __init__(self, enable_timing): pass def record(self): self.time = time.perf_counter() def elapsed_time(self, end_event): assert isinstance(end_event, Event) return end_event.time - self.time def gen_sparse_csr(shape, nnz): fill_value = 0 total_values = functools.reduce(operator.mul, shape, 1) dense = np.random.randn(total_values) fills = random.sample(list(range(total_values)), total_values - nnz) for f in fills: dense[f] = fill_value dense = torch.from_numpy(dense.reshape(shape)) return dense.to_sparse_csr() def gen_sparse_coo(shape, nnz): dense = np.random.randn(*shape) values = [] indices = [[], []] for n in range(nnz): row = random.randint(0, shape[0] - 1) col = random.randint(0, shape[1] - 1) indices[0].append(row) indices[1].append(col) values.append(dense[row, col]) return torch.sparse_coo_tensor(indices, values, size=shape) def gen_sparse_coo_and_csr(shape, nnz): total_values = functools.reduce(operator.mul, shape, 1) dense = np.random.randn(total_values) fills = random.sample(list(range(total_values)), total_values - nnz) for f in fills: dense[f] = 0 dense = torch.from_numpy(dense.reshape(shape)) return dense.to_sparse(), dense.to_sparse_csr()