How to add a sparse tensor to a dense parameter quickly

username: Bert84258836
cooked: <p>I am doing some benchmarking using EmbeddingBag vs using a nn.linear layer with torch sparse tensors. My data is bag of words tfidf features which are sparse (N x 16 million) where N is number of documents. The two methods I am comparing are</p>
<pre><code class="lang-auto">class EmbedLR(nn.Module):
    def __init__(self, vocab_size, alpha, learning_rate):
        super(EmbedLR, self).__init__()
        self.embed = nn.EmbeddingBag(
            num_embeddings=vocab_size,
            embedding_dim=1,
            mode="sum",
            sparse=True
        )

        self.learning_rate = learning_rate

        initial_eta0 = np.sqrt(1.0 / np.sqrt(alpha))
        self.optimal_init = 1.0 / (initial_eta0 * alpha)
        self.optimizer = optim.SGD(self.parameters(), self.learning_rate)

    # gets csr form for input to embedding
    def compute_offsets(self, batch):
        return torch.LongTensor(batch.indices), torch.LongTensor(batch.indptr[:-1]), torch.FloatTensor(batch.data)

    def forward(self, x):
        input, offsets, per_sample_weights = self.compute_offsets(x)
        return self.embed(
            input=input,
            offsets=offsets,
            per_sample_weights=per_sample_weights
        )

    def zero_init(self):
        nn.init.zeros_(self.embed.weight)

class LogisticRegression(nn.Module):
    def __init__(self, vocab_size, alpha, learning_rate):
        super(LogisticRegression, self).__init__()

        self.vocab_size = vocab_size
        self.alpha = alpha
        self.learning_rate = learning_rate

        self.linear = nn.Linear(self.vocab_size, 1, bias=False)

        initial_eta0 = np.sqrt(1.0 / np.sqrt(self.alpha))
        self.optimal_init = 1.0 / (initial_eta0 * self.alpha)
        self.optimizer = optim.SGD(self.parameters(), self.learning_rate)

   # x is a sparse torch tensor
    def forward(self, x):
        return self.linear(x)

    def zero_init(self):
        nn.init.zeros_(self.linear.weight)
</code></pre>
<p>The output of running these models on the same data is as follows:</p>
<p>LogisticRegression:</p>
<p>Forward time: 0.6464931964874268<br>
Loss value: 0.6931458711624146<br>
Loss backward time: 0.671107292175293<br>
p.add_(d_p, alpha=-group[lr]) takes 0.0021352767944335938<br>
Optimizer step time: 0.0022284984588623047<br>
Optimizer zero grad time: 0.0012955665588378906<br>
Converge in 1.3219084739685059 seconds<br>
Convergence per doc: 7.284446321532517e-05</p>
<p>EmbedLR:</p>
<p>Forward time: 0.029858112335205078<br>
Loss value: 0.6931458711624146<br>
Loss backward time: 0.0931253433227539<br>
p.add_(d_p, alpha=-group[lr]) takes 3.497199296951294<br>
Optimizer step time: 3.497309923171997<br>
Optimizer zero grad time: 0.004025459289550781<br>
Converge in 3.624892234802246 seconds<br>
Convergence per doc: 0.00019975159722280522</p>
<p>Using EmbeddingBag the forward pass is faster than using a linear layer on a sparse tensor.  But the backward pass is slower. The main bottleneck is the time it takes for the optimizer step. Specifically within the SGD optimizer this line of code is much faster<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/cd9b27231b51633e76e28b6a34002ab83b0660fc/torch/optim/sgd.py#L98" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/cd9b27231b51633e76e28b6a34002ab83b0660fc/torch/optim/sgd.py#L98" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/cd9b27231b51633e76e28b6a34002ab83b0660fc/torch/optim/sgd.py#L98</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="88" style="counter-reset: li-counter 87 ;">
<li>            if 'momentum_buffer' not in param_state:</li>
<li>                buf = param_state['momentum_buffer'] = d_p.clone()</li>
<li>            else:</li>
<li>                buf = param_state['momentum_buffer']</li>
<li>                buf.mul_(momentum).add_(1 - dampening, d_p)</li>
<li>            if nesterov:</li>
<li>                d_p = d_p.add(momentum, buf)</li>
<li>            else:</li>
<li>                d_p = buf</li>
<li>
</li>
<li class="selected">        p.data.add_(-group['lr'], d_p)</li>
<li>
</li>
<li>return loss</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>p is the parameter and d_p is the gradient of the parameter (in this case weight vector either linear.weight or embed.weight). In the EmbedLR d_p is a sparse tensor and p is dense tensor and I believe this is why it is slow (dense-sparse addition). It is faster in LogisticRegression as both p and d_p are dense tensors.</p>
<p>I am wondering how I can speed this optimizer step bottle neck up. Is there a good way of doing dense sparse addition. I tried making the parameter p itself sparse like:</p>
<p><code>self.embed.weight = nn.Parameter(torch.sparse.FloatTensor(16777216, 1))</code></p>
<p>but got the following error:</p>
<p><code>RuntimeError: Could not run 'aten::_embedding_bag' with arguments from the 'SparseCPUTensorId' backend. 'aten::_embedding_bag' is only available for these backends: [CPUTensorId, CUDATensorId, VariableTensorId].</code></p>