Sparse Layers for Sequential NN Models

Dan_Bingo: <p>I’m looking for a method to sparsify a simple network as described below:</p>
<pre><code class="lang-auto">    model = torch.nn.Sequential(
        collections.OrderedDict(
            [
                ("layer1", torch.nn.Linear(num_A, num_A)),
                ("act1", torch.nn.Tanh()), 
                ("layer2", torch.nn.Linear(num_A, num_B)),
                ("act2", torch.nn.Tanh()), 
                ("layer3", torch.nn.Linear(num_B, num_B)),
            ]
        )
    )
</code></pre>
<p>I am using the <code>torch.nn.utils.prune.custom_from_mask</code> to prune the weights I want to be zero by sending a matrix to the device that is 99% zeros with 1% ones.</p>
<pre><code class="lang-auto">matrix = matrix.to(device)
module1 = model.layer1
module1 = module1.to(device)
torch.nn.utils.prune.custom_from_mask(module1, name='weight', mask=matrix)
</code></pre>
<p>The model size after masking is still as large as a fully connected network. I believe that the results I’m getting are good for the model and this strategy, but I need to make the network sparse to deploy this at scale due to memory limitations of making the matrix dense (on the CPU &amp; GPU) and building the layer fully connected (on the GPU).</p>
<p>Any suggestions would be appreciated.</p>

KFrank: <p>Hi Dan!</p>
<aside class="quote no-group" data-username="Dan_Bingo" data-post="1" data-topic="159718" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/35a633/48.png" class="avatar"> Dan_Bingo:</div>
<blockquote>
<p>I’m looking for a method to sparsify a simple network</p>
</blockquote>
</aside>
<p>If you intend to use your network just for inference (e.g., after deployment<br>
to a smaller device), you could do something like:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; lin = torch.nn.Linear (3, 5)
&gt;&gt;&gt; mask = torch.randint (2, (5, 3))
&gt;&gt;&gt;
&gt;&gt;&gt; with torch.no_grad():
...     lin.weight.mul_ (mask)   # weight has a lot of zeros but not in sparse format
...
Parameter containing:
tensor([[-0.0000,  0.4872,  0.0000],
        [-0.1372, -0.0000, -0.0000],
        [ 0.0000,  0.3245, -0.0000],
        [ 0.0000,  0.0000, -0.3115],
        [ 0.0000,  0.2717,  0.0000]], requires_grad=True)
&gt;&gt;&gt;
&gt;&gt;&gt; input = torch.randn (3)
&gt;&gt;&gt;
&gt;&gt;&gt; output = lin (input)
&gt;&gt;&gt;
&gt;&gt;&gt; lin.weight = torch.nn.Parameter (lin.weight.to_sparse())   # weight is now in sparse format
&gt;&gt;&gt; lin.weight
Parameter containing:
tensor(indices=tensor([[0, 1, 2, 3, 4],
                       [1, 0, 1, 2, 1]]),
       values=tensor([ 0.4872, -0.1372,  0.3245, -0.3115,  0.2717]),
       size=(5, 3), nnz=5, layout=torch.sparse_coo, requires_grad=True)
&gt;&gt;&gt;
&gt;&gt;&gt; outputB = lin (input)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (output, outputB)   # same result
True
</code></pre>
<p>Note, you can’t use this for training as pytorch does not yet support<br>
sparse-tensor backpropagation.</p>
<p>Best.</p>
<p>K. Frank</p>

Dan_Bingo: <p>Thank you, K. Frank. I think this will do nicely for deploying the model.</p>
<p>I do need a way to train dense-sparse and sparse-sparse layers.</p>
<p>I’m taking a look at the RigL PyTorch implementation here: <a href="https://pypi.org/project/rigl-torch/" class="inline-onebox" rel="noopener nofollow ugc">rigl-torch · PyPI</a></p>

Dan_Bingo: <p>I was able to modify the rigl-torch implementation to use a static sparse tensor to prune the weights instead of relying on the random selection from the uniform distribution as done in the paper. Comparing results from the torch.nn.prune.custom_from_mask implementation, the training time is much slower with RigL than with the masking, but the model size and inference times are multiples faster with RigL.</p>