文章目录
- 0. 前言
- 1. `torch.nn.utils.prune`中的结构化剪枝方法
- 2. PyTorch实例
- 2.1 `random_structured`
- 2.2 `prune.ln_structured`
- 3. 总结
0. 前言
按照国际惯例，首先声明：本文只是我自己学习的理解，虽然参考了他人的宝贵见解及成果，但是内容可能存在不准确的地方。如果发现文中错误，希望批评指正，共同进步。
在前文：【PyTorch单点知识】神经元网络模型剪枝prune模块介绍（上，非结构化剪枝）中介绍了PyTorch中的prune模型剪枝模块中的非结构化剪枝。本文将通过实例说明utils.prune中的结构化剪枝方法。
1. torch.nn.utils.prune中的结构化剪枝方法
2. PyTorch实例
首先建立一个简单的模型：
python">import torch
import torch.nn as nn
from torch.nn.utils import prunetorch.manual_seed(888)
class SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()# 创建一个简单的卷积层self.conv = nn.Conv2d(in_channels=1, out_channels=3, kernel_size=3)model = SimpleModel()
通过print(model.conv.weight)可以打印出权重为：
python">Parameter containing:
tensor([[[[-0.3017, 0.1290, -0.2468],[ 0.2107, 0.1799, 0.1923],[ 0.1887, -0.0527, 0.1403]]],[[[ 0.0799, 0.1399, -0.0084],[ 0.2013, -0.0352, -0.1027],[-0.1724, -0.3094, -0.2382]]],[[[ 0.0419, 0.2224, -0.1558],[ 0.2084, 0.0543, 0.0647],[ 0.1493, 0.2011, 0.0310]]]], requires_grad=True)
2.1 random_structured
这个方法会在指定维度dim（默认为-1）上剪枝一个随机通道：
python">prune.random_structured(model.conv, name="weight", amount=0.33)
print("Weight after RandomStructured pruning (33%):")
print(model.conv.weight)
输出为：
python">Weight after RandomStructured pruning (33%):
tensor([[[[-0.3017, 0.1290, -0.0000],[ 0.2107, 0.1799, 0.0000],[ 0.1887, -0.0527, 0.0000]]],[[[ 0.0799, 0.1399, -0.0000],[ 0.2013, -0.0352, -0.0000],[-0.1724, -0.3094, -0.0000]]],[[[ 0.0419, 0.2224, -0.0000],[ 0.2084, 0.0543, 0.0000],[ 0.1493, 0.2011, 0.0000]]]], grad_fn=<MulBackward0>)
由于权重的维度为[3, 1, 3, 3]，我们也可以试试在其他维度（dim=0和dim=2）上进行剪枝：
dim=0
python">prune.random_structured(model.conv, name="weight", amount=0.33,dim=0)
print(model.conv.weight)
输出为：
python">tensor([[[[-0.3017, 0.1290, -0.2468],[ 0.2107, 0.1799, 0.1923],[ 0.1887, -0.0527, 0.1403]]],[[[ 0.0799, 0.1399, -0.0084],[ 0.2013, -0.0352, -0.1027],[-0.1724, -0.3094, -0.2382]]],[[[ 0.0000, 0.0000, -0.0000],[ 0.0000, 0.0000, 0.0000],[ 0.0000, 0.0000, 0.0000]]]], grad_fn=<MulBackward0>)
dim=2
python">prune.random_structured(model.conv, name="weight", amount=0.33,dim=2)
print(model.conv.weight)
输出为：
python">tensor([[[[-0.3017, 0.1290, -0.2468],[ 0.2107, 0.1799, 0.1923],[ 0.0000, -0.0000, 0.0000]]],[[[ 0.0799, 0.1399, -0.0084],[ 0.2013, -0.0352, -0.1027],[-0.0000, -0.0000, -0.0000]]],[[[ 0.0419, 0.2224, -0.1558],[ 0.2084, 0.0543, 0.0647],[ 0.0000, 0.0000, 0.0000]]]], grad_fn=<MulBackward0>)
2.2 prune.ln_structured
这个方法会在指定维度dim（默认为-1）上按最小n范数剪枝一个通道：
python">prune.ln_structured(model.conv, name="weight",amount=0.33,n=1,dim=-1)
print(model.conv.weight)
输出为：
python">tensor([[[[-0.3017, 0.1290, -0.0000],[ 0.2107, 0.1799, 0.0000],[ 0.1887, -0.0527, 0.0000]]],[[[ 0.0799, 0.1399, -0.0000],[ 0.2013, -0.0352, -0.0000],[-0.1724, -0.3094, -0.0000]]],[[[ 0.0419, 0.2224, -0.0000],[ 0.2084, 0.0543, 0.0000],[ 0.1493, 0.2011, 0.0000]]]], grad_fn=<MulBackward0>)
更改dim也是同样的效果：
dim=0
python">prune.ln_structured(model.conv, name="weight",amount=0.33,n=1,dim=0)
print(model.conv.weight)
输出为：
python">tensor([[[[-0.3017, 0.1290, -0.2468],[ 0.2107, 0.1799, 0.1923],[ 0.1887, -0.0527, 0.1403]]],[[[ 0.0799, 0.1399, -0.0084],[ 0.2013, -0.0352, -0.1027],[-0.1724, -0.3094, -0.2382]]],[[[ 0.0000, 0.0000, -0.0000],[ 0.0000, 0.0000, 0.0000],[ 0.0000, 0.0000, 0.0000]]]], grad_fn=<MulBackward0>)
dim=2
python">prune.ln_structured(model.conv, name="weight",amount=0.33,n=1,dim=2)
print(model.conv.weight)
输出为：
python">tensor([[[[-0.3017, 0.1290, -0.2468],[ 0.0000, 0.0000, 0.0000],[ 0.1887, -0.0527, 0.1403]]],[[[ 0.0799, 0.1399, -0.0084],[ 0.0000, -0.0000, -0.0000],[-0.1724, -0.3094, -0.2382]]],[[[ 0.0419, 0.2224, -0.1558],[ 0.0000, 0.0000, 0.0000],[ 0.1493, 0.2011, 0.0310]]]], grad_fn=<MulBackward0>)
