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

1

I want to do forward pass (inference) with a sparse PyTorch model, but I got the error 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.

I wonder whether PyTorch supports doing the forward pass between a sparse model and a (normal) dense input image? Is this even possible?

=====================================
Below is the detailed context.

I followed the tutorial of PyTorch pruning, I am aware that torch.prune 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 tutorial of PyTorch sparse to convert the pruned model into a sparse model.

Below shows how I convert the pruned model into a sparse model:

layer.weight = torch.nn.Parameter(layer.weight.to_sparse())
layer.bias = torch.nn.Parameter(layer.bias.to_sparse())

The complete code to reproduce the error is shown below:

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)

Thanks a lot for the clarification!

2

I read and learn a lot from your comments @ptrblck, maybe you could help me?

Thank you!

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.