I’d like to keep keep in the tensor only K largest elements in each row (corresponding to logits/logprobs of K highest scoring classes) to optimize disk space during serialization. I’d like to get a sparse tensor as an output.
Is there an simpler way than this (e.g. directly passing indices from topk to torch.sparse.FloatTensor constructor) ?
x = torch.rand(128, 512, 40) # 128 and 512 are batch dimension, 40 is class logits dimension

def sparse_topk(x, K, dim = -1):
  return torch.zeros_like(x).scatter_(dim, x.topk(max(1, min(K, x.size(dim))), dim = dim).indices, 1.0).to_sparse()

torch.save(x, 'x.pt') # 11M     x.pt
torch.save(sparse_topk(x, K = 5), 'x_.pt') # 8.8M    x_.pt
torch.save(sparse_topk(x.flatten(end_dim = 1), K = 5), 'x__.pt') # 6.3M    x__.pt

I’m a bit confused that the file size depends on if I flatten the batch dimensions or not. Should I keep the sparse dimension first (and thus transpose my tensor)?
Ideally, I’d also like to store indices as a bitset (since my number of classes is smaller than 256).
Any other advice on saving disk space?
Thanks!