How can I apply a linear transformation on sparse matrix in PyTorch?

Maryam_Khaliji: <p>In PyTorch, we have <code>nn.linear</code> that applies a linear transformation to the incoming data:</p>
<pre><code class="lang-auto">y = WA+b
</code></pre>
<p>In this formula, <code>W</code> and <code>b</code> are our learnable parameters and <code>A</code> is my input data matrix. The matrix <code>A</code> for my case is too large for RAM to complete loading, so I use it sparsely. Is it possible to perform such an operation on sparse matrices using PyTorch?</p>

Karthik_Ganesan: <p>When you say you use it ‘sparsely’, do you mean you are using sparse tensors (i.e., tensors where a large proportion of the values are 0)? If so, then you can just represent your tensors using <a href="https://pytorch.org/docs/stable/sparse.html" rel="noopener nofollow ugc"><code>torch.sparse</code></a> to do this. And then you can use <code>torch.addmm</code> to implement the linear layer.</p>

Maryam_Khaliji: <aside class="quote no-group" data-username="Karthik_Ganesan" data-post="2" data-topic="160139">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/karthik_ganesan/48/34393_2.png" class="avatar"> Karthik_Ganesan:</div>
<blockquote>
<p>torch.addmm</p>
</blockquote>
</aside>
<p>Thanks for your response. Yes, I use COO format.<br>
When I define <code>b' as bias and </code>W<code> as weight matrix, I use nn.parameter(). When we just have two learnable parameters (</code>W<code>and</code>b`), and the COO sparse matrix is our input data (not learnable parameter), which of torch.addmm() and torch.sparse.addmm() should I use?</p>

KFrank: <p>Hi Maryam!</p>
<aside class="quote no-group quote-modified" data-username="Maryam_Khaliji" data-post="3" data-topic="160139" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/maryam_khaliji/48/45054_2.png" class="avatar"> Maryam_Khaliji:</div>
<blockquote>
<p>I define <code>b' as bias and </code>W` as weight matrix, I use nn.parameter().</p>
</blockquote>
</aside>
<p>Note, you can use a <code>Linear</code> and apply it to a sparse tensor.  (You don’t<br>
need to explicitly call something like <code>torch.sparse.addmm()</code>.)</p>
<p>You don’t have to use <code>Linear</code>, of course, but it might be stylistically<br>
better since it’s used so commonly.</p>
<p>If it better fits your use case, you can also use the functional form,<br>
<code>torch.nn.functional.linear()</code>.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt; s = ((torch.randn (3, 5) &gt; 0.0) * torch.randn (3, 5)).to_sparse()   # make a sparse tensor
&gt;&gt;&gt; s
tensor(indices=tensor([[0, 0, 0, 0, 1, 1, 2, 2, 2],
                       [0, 1, 2, 3, 2, 3, 2, 3, 4]]),
       values=tensor([-1.0556, -0.5344,  1.2262,  1.1603, -0.1418, -0.3402,
                       0.6979, -0.5543, -0.2561]),
       size=(3, 5), nnz=9, layout=torch.sparse_coo)
&gt;&gt;&gt; lin = torch.nn.Linear (5, 3)   # make a (non-sparse) Linear
&gt;&gt;&gt; t = lin (s)
&gt;&gt;&gt; t
tensor([[-0.4231, -0.6271, -0.6817],
        [-0.2014, -0.0587,  0.1683],
        [-0.3242,  0.1646,  0.3315]], grad_fn=&lt;AddmmBackward0&gt;)
&gt;&gt;&gt; t.sum().backward()   # backward works
&gt;&gt;&gt; lin.weight.grad
tensor([[-1.0556, -0.5344,  1.7823,  0.2658, -0.2561],
        [-1.0556, -0.5344,  1.7823,  0.2658, -0.2561],
        [-1.0556, -0.5344,  1.7823,  0.2658, -0.2561]])
&gt;&gt;&gt; w = torch.nn.Parameter (lin.weight.detach().clone())
&gt;&gt;&gt; b = torch.nn.Parameter (lin.bias.detach().clone())
&gt;&gt;&gt; t = torch.nn.functional.linear (s, w, b)   # functional form of linear also works
&gt;&gt;&gt; t.sum().backward()
&gt;&gt;&gt; w.grad
tensor([[-1.0556, -0.5344,  1.7823,  0.2658, -0.2561],
        [-1.0556, -0.5344,  1.7823,  0.2658, -0.2561],
        [-1.0556, -0.5344,  1.7823,  0.2658, -0.2561]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Maryam_Khaliji: <p>Thank you so much. It was very helpful.</p>

Jeremie_Kalfon: <p>Hello,</p>
<p>KFrank’s solution doesn’t work on my end. I get that torch sparse tensor don’t have a reshape.(RuntimeError: reshape is not implemented for sparse tensors)</p>
<p>I think that only addmm would work…<br>
I guess doing something like this?</p>
<pre><code class="lang-auto">class SparseLinear(nn.Module):
    def __init__(self, insize, outsize):
        super(SparseLinear, self)
        self.weights = nn.Parameter(torch.randn(insize, outsize))
        self.bias = nn.Parameter(torch.randn(outsize))

    def forward(self, x):
        return torch.sparse.addmm(self.bias, x, self.weights)
</code></pre>