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

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

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

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

修剪后使用网络进行推断，这将首先关闭一些与掩码的连接，然后运行推理。
将原始权重归零，然后从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。模块，只有张量对象。这里有一些有用的链接：
https://github.com/pytorch/pytorch/issues/36214#issuecomment-619586452
https://pytorch.org/docs/stable/sparse.html

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

结构化修剪：通过删除张量的整行/列来减小权重张量的尺寸。这转化为去除神经元及其所有传入和传出连接(在密集层中(或整个卷积过滤器(在卷积层中(。
非结构化修剪：单个权重可以"删除"(归零(，而不受最终张量形状的限制。这意味着删除神经元之间的单个连接(在密集层中(或删除卷积过滤器的单个权重(在卷积层中(。请注意，生成的权重张量可以是稀疏的，但保持其原始形状。

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

非结构化修剪只有在去除大量权重时才能帮助提高性能。在这种情况下，您可以依赖 PyTorch 稀疏操作，也可以尝试查找包含所有零的行/列，因此可以删除。

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

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

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

它将减少多达 85% 的时间而不会降低精度。

通过将权重设置为 0.0 来修剪权重只是故事的一半。另一半是将它们从模型中(物理上(删除，以便它们不参与任何计算。这在非结构化修剪中并非易事，但是对于结构化修剪，如果您能够从 conv 操作中删除特定的输出通道(从而删除下一个 conv 的输入通道(，那么您可以在模型中获得一些加速。

要物理删除通道，一旦确定了要删除的通道，就需要调整 Conv2d 模块上的可学习权重参数。您可能还需要更改它之后的几个图层，特别是批处理范数层和下一个卷积或线性层。

相关内容

最新更新

热门标签：