修剪模型不会提高推理速度或减小模型大小

Question

我正在尝试使用torch.nn.utils.prune提供 2 个张量的PyTorch 修剪我的模型，

1. 一个是原始重量和
2. 另一个是包含 0 和 1 的掩码，帮助我们关闭网络中的某些连接。

我已经尝试了两种解决方案，但都没有提高推理速度：

1. 使用剪枝后的网络进行推断，这将首先关闭一些与掩码的连接，然后运行推断。
2. 使用掩码将原始权重归零，然后从 state_dict 中删除掩码以进行推断。

有没有办法用模型张量和掩码来提高速度？与 0 的非零浮点数相乘不会比将 2 个浮点数相乘快吗？
这是我的修剪函数和修剪速度计算程序：

def prune_net(net):
    """Prune 20% net's weights that have abs(value) approx. 0
    Function that will be use when an iteration is reach
    Args:

    Return:
        newnet (nn.Module): a newnet contain mask that help prune network's weight
    """
    if not isinstance(net,nn.Module):
        print('Invalid input. Must be nn.Module')
        return
    newnet = copy.copy(net)
    modules_list = []

    for name, module in newnet.named_modules():
        if isinstance(module, torch.nn.Conv2d):
            modules_list += [(module,'weight'),(module,'bias')]
        if isinstance(module, torch.nn.Linear):
            modules_list += [(module,'weight'),(module,'bias')]

    prune.global_unstructured(
        modules_list,
        pruning_method=prune.L1Unstructured,
        amount=0.2,)
    return newnet

测试推理速度第一种情况：

import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F
import time
from torch.autograd import Variable


torch.set_default_tensor_type('torch.cuda.FloatTensor')
old_net = init_your_net()

new_net = prune_net(old_net)
new_net = prune_net(new_net)

old_net.eval()
new_net.eval()

old_net = old_net.cuda()
new_net = new_net.cuda()
dataset = load_your_dataset()

for i in range(100):
    x = dataset[i]
    x = x.cuda()
    y = x.cuda()

    #new infer
    start_time = time.perf_counter()
    detections = new_net(x).data
    time_new += time.perf_counter() - start_time

    #old infer
    start_time = time.perf_counter()
    detections = old_net(y).data
    time_old += time.perf_counter() - start_time
print('old ',time_old)
print('new ', time_new)

测试推理速度第二种情况：

import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F
import time
from torch.autograd import Variable


torch.set_default_tensor_type('torch.cuda.FloatTensor')
old_net = init_your_net()

new_net = prune_net(old_net)
new_net = prune_net(new_net)
# Apply mask to model tensor and remove mask from state_dict
for name, module in new_net.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.remove(module,'weight')
        prune.remove(module,'bias')
    if isinstance(module, torch.nn.Linear):
        prune.remove(module,'weight')
        prune.remove(module,'bias')

old_net.eval()
new_net.eval()

old_net = old_net.cuda()
new_net = new_net.cuda()
dataset = load_your_dataset()

for i in range(100):
    x = dataset[i]
    x = x.cuda()
    y = x.cuda()

    #new infer
    start_time = time.perf_counter()
    detections = new_net(x).data
    time_new += time.perf_counter() - start_time

    #old infer
    start_time = time.perf_counter()
    detections = old_net(y).data
    time_old += time.perf_counter() - start_time
print('old ',time_old)
print('new ', time_new)

更新
我发现火炬有一个稀疏模块，如果我们修剪足够的参数，它可以减少内存使用，但它还不支持 nn.Module，只支持 Tensor 对象。这里有一些有用的链接：
https : //github.com/pytorch/pytorch/issues/36214#issuecomment-619586452
https://pytorch.org/docs/stable/sparse.html

Answer 1

了解非结构化修剪和结构化修剪之间的区别很重要。

• 结构化修剪：通过移除张量的整行/列来减少权重张量的维度。这转化为移除神经元及其所有传入和传出连接（在密集层中）或整个卷积滤波器（在卷积层中）。

• 非结构化修剪：可以“删除”（清零）单个权重，而不受最终张量形状的限制。这转化为移除神经元之间的单个连接（在密集层中）或移除卷积滤波器的单个权重（在卷积层中）。请注意，生成的权重张量可以是稀疏的，但会保持其原始形状。

目前torch.nn.utils.prune只支持非结构化剪枝，几乎没有帮助降低推理成本，因为 GPU 没有针对稀疏矩阵乘法进行优化。虽然您可能希望减少权重张量的维度以减少浮点运算的数量，但非结构化修剪会生成具有许多零的权重张量，但不会自动减小此类张量的大小。

只有在移除大量权重时，非结构化修剪才能帮助提高性能。在这种情况下，您可以依靠PyTorch 稀疏操作或尝试查找包含全零的行/列，从而可以将其删除。

相反，如果您想研究结构化修剪，您可以查看TorchPruner，这是我为研究目的自己开发的一个库，它提供实用程序来查找最不重要的神经元并相应地对权重张量进行切片。

Answer 2

我也在尝试修剪以提高推理速度。但我发现更有用的是使用 ONNX 和 ONNXRuntime。这是包含所有步骤的链接：

https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

它将在不损失精度的情况下减少高达 85% 的时间。