nn.init.sparse poor performance #6021
Description
- OS: Ubuntu 16.04
- PyTorch version: 0.3.1
- How you installed PyTorch (conda, pip, source): conda
- Python version: 3.6.4
- CUDA/cuDNN version: 9.1
- GPU models and configuration: GTX 1060
It seems that nn.init.sparse takes too much time for its purpose. On my CPU, a 500x500 tensor initialization takes 0.25 seconds.
I wrote a piece of code that can be 10x faster than nn.init.sparse. I would ask you to try it yourself on your machines and check whether the semantic is the same.
import torch
import torch.nn as nn
import time
sparsity = 0.9
tensor_size = 500
t1 = torch.Tensor(tensor_size, tensor_size)
s = time.time()
nn.init.sparse(t1, sparsity)
print(time.time() - s)
t2 = torch.Tensor(tensor_size * tensor_size)
s = time.time()
t2.normal_()
if sparsity > 0:
zero_weights = torch.randperm(int(tensor_size * tensor_size))
zero_weights = zero_weights[:round(tensor_size * tensor_size * sparsity)]
t2[zero_weights] = 0
t2 = t2.view(tensor_size, tensor_size)
print(time.time() - s)