Forward pass between a sparse PyTorch model and a dense input image

SunHaozhe: <p>I want to do forward pass (inference) with a sparse PyTorch model, but I got the error <code>RuntimeError: Input type (torch.FloatTensor) and weight type (torch.sparse.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor</code>.</p>
<p>I wonder whether PyTorch supports doing the forward pass between a sparse model and a (normal) dense input image? Is this even possible?</p>
<p>=====================================<br>
Below is the detailed context.</p>
<p>I followed the <a href="https://pytorch.org/tutorials/intermediate/pruning_tutorial.html" rel="noopener nofollow ugc">tutorial of PyTorch pruning</a>, I am aware that <code>torch.prune</code> merely sets entries to 0, it does not actually make the model more sparse in the physical sense. That’s why I then followed the <a href="https://pytorch.org/docs/stable/sparse.html" rel="noopener nofollow ugc">tutorial of PyTorch sparse</a> to convert the pruned model into a sparse model.</p>
<p>Below shows how I convert the pruned model into a sparse model:</p>
<pre><code class="lang-auto">layer.weight = torch.nn.Parameter(layer.weight.to_sparse())
layer.bias = torch.nn.Parameter(layer.bias.to_sparse())
</code></pre>
<p>The complete code to reproduce the error is shown below:</p>
<pre><code class="lang-auto">import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# a simple LeNet model as an example
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet()

dummy_input = torch.randn((2, 1, 28, 28))
dummy_output = model(dummy_input)
print(dummy_output.shape)

modules_to_prune = [
    model.conv1,
    model.conv2,
    model.fc1,
    model.fc2,
    model.fc3
]

for layer in modules_to_prune:
    parameters_to_prune = [(layer, "weight"), (layer, "bias")]

    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.RandomUnstructured,
        amount=0.2,
    )

    prune.remove(layer, 'weight')
    prune.remove(layer, 'bias')
    
    layer.weight = torch.nn.Parameter(layer.weight.to_sparse())
    layer.bias = torch.nn.Parameter(layer.bias.to_sparse())

dummy_input = torch.randn((2, 1, 28, 28))
dummy_output = model(dummy_input)
print(dummy_output.shape)

</code></pre>
<p>Thanks a lot for the clarification!</p>

SunHaozhe: <p>I read and learn a lot from your comments <a class="mention" href="/u/ptrblck">@ptrblck</a>, maybe you could help me?</p>
<p>Thank you!</p>
<p>My use case is unstructured pruning of neural networks. I want the pruned model to be effectively sparse (pruned model should occupy less physical space, inference time may also be effectively reduced), not just setting pruned parameters to zeros.</p>