A unified framework for Multi-Modality Structured Pruning
torch
torchvision
numpy
tqdm
thop
- (optional)
pynvml
pyRAPL
jtop
torch2onnx
deform_conv2d_onnx_exporter
# official
pip install -U Unified-Pruning
# nightly (Recommendation)
pip install git+https://github.com/Nobreakfast/UniP.git
# development version
git clone https://github.com/Nobreakfast/UniP.git
cd UniP
pip install -e .
import torch
import torchvision.models as models
from unip.core.pruner import BasePruner
from unip.utils.evaluation import cal_flops
# load model and example input
model = models.resnet18()
example_input = torch.rand(1, 3, 224, 224, requires_grad=True)
# record the flops and params
cal_flops(model, example_input, device)
# define pruner
pruner = BasePruner(
model,
example_input,
"UniformRatio",
algo_args={"score_fn": "weight_sum_l1_out"},
)
pruner.prune(0.3)
# record the flops and params
cal_flops(model, example_input, device)
import sys
sys.path.append("./tests")
import torch
from model.radarnet import RCNet
from model.backbone.conv_utils.dcn import DeformableConv2d
from unip.core.pruner import BasePruner
from unip.utils.evaluation import cal_flops
# RCNet need more customized node for deformable conv
from unip.core.node import dcnNode
# load model and example input
model = RCNet(in_channels=3)
example_input = torch.randn(1, 3, 320, 320, requires_grad=True)
# record the flops and params
cal_flops(model, example_input)
# define a dict indicate the module with node
igtype2nodetype = {DeformableConv2d: dcnNode}
# define pruner
pruner = BasePruner(
model,
example_input,
"UniformRatio",
algo_args={"score_fn": "weight_sum_l1_out"},
igtype2nodetype=igtype2nodetype,
)
pruner.prune(0.3)
# record the flops and params
cal_flops(model, example_input, device)
Please refer to the ./tutorials
folder for more examples.
Before report bugs by opening an issue in the GitHub Issue Tracker, you may follow the instructions of test_Unip README.md. Basically, you need to provide the following information:
- the model you want to prune
- an example inference Once you new a issue, we will fix it as soon as possible.
You are welcome to contribute to this project. Please follow the Contribution Guide.
- Nodes for operations:
BaseNode
: Base class for all other nodesInOutNode
: Node for conv, linear, and bundle parametersOutOutNode
: Node for bn, ln, gn, and dwconvInInNode
: Node for add, sub, mul, and divRemapNode
: Node for concat, split, and sliceChangeNode
: Node for reshape, permute, expand, transpose, matmul, and upsampleDummyNode
: Node for dummy input and outputActivationNode
: Node for activationPoolNode
: Node for adaptive pooling, max pooling, and average poolingCustomNode
: Node for custom module
- Algorithm:
BaseAlgo
: Base class for all other algorithmsRandomAlgo
: Random ratio and indexUniformAlgo
: Uniform ratio and random index
- Tested models:
- torchvision:
- GoogleNet
- VGG
- ResNet, WideResNet,
- MobileNetV2, MobileNetV3
- EfficientNet, EfficientNetV2
- SqueezeNet
- mobilevit
- Achelous
- torchvision:
- Tested modules:
- conv: basic conv, depth-wise conv, deformable conv (with Custom Node
dcnNode
) - linear: basic linear, linear with input more than 2 dimensions
- calculation operators: concat, split, slice, reshape, permute, expand, transpose, matmul, upsample, add, sub, div, mul, squeeze, unsqueeze
- fired design: residual, ghost module, attention, shuffle attention
- conv: basic conv, depth-wise conv, deformable conv (with Custom Node
- support backwards type
- ConvolutionBackward0
- AddmmBackward0, MmBackward0, BmmBackward0
- AddBackward0, SubBackward0, MulBackward0, DivBackward0
- NativeBatchNormBackward0, NativeLayerNormBackward0, NativeGroupNormBackward0
- CatBackward0, SplitBackward0, SliceBackward0
- ReshapeAliasBackward0, UnsafeViewBackward0, ViewBackward0, PermuteBackward0, ExpandBackward0, TransposeBackward0, SqueezeBackward1, UnsqueezeBackward0,
- MeanBackward1, MaxPool2DWithIndicesBackward0, AvgPool2DBackward0, UpsampleBilinear2DBackward0, UpsampleNearest2DBackward0, UpsampleBicubic2DBackward0
- ReluBackward0, SiluBackward0, GeluBackward0, HardswishBackward0, SigmoidBackward0, TanhBackward0, SoftmaxBackward0, LogSoftmaxBackward0
- AccumulateGrad, TBackward0, CloneBackward0
- RepeatBackward0
- EmbeddingBackward0
- AsStridedBackward0
- CopySlices
- UnsafeSplitBackward0
- UnbindBackward0
- IndexBackward0
- SqueezeBackward0
- MaxBackward0
- UnsafeSplitBackward0
- StackBackward0
- TransposeBackward1
- when
in_channels
greater thangroups
- when operation
conv
andfc
does not usePyTorch
module implementation - fix the bug of useless
split
, when the prev_node's group isnon-prunable
-
TransposeConv
error - for some nodes starting from a non-
Input
, the dim_offset is wrong -
ConcatNode
is the next ofReshapeNode
-
nn.MultiheadAttention
module not working: fixed by addingCustomNode
- RCNet may failed cuz the residual with input
- fix the bug for
DCN
module: usedcnNode
- fix the bug for such module like GhostModule, use Non-
InOutNode
beforeOutputNode
- does not prune the
LastLienarNode
forto_qkv
like module - need to fix the
round_to
like attribute forto_qkv
like module -
dim_offset
for reshape node is not always correct
pynvml
: 11.5.0pyRAPL
: 0.2.3.1
Note: require to give read permission to the specific file:
sudo chmod -R a+r /sys/class/powercap/intel-rapl
import torch
import torchvision.models as models
from unip.utils.energy import Calculator
# define a device dict
device_dict = {
"NvidiaDev": {'device_id': 0},
"IntelDev": {},
}
calculator = Calculator(device_dict)
model = models.resnet18()
model.eval().cuda()
example_input = torch.randn(1, 3, 224, 224).cuda()
@calculator.measure(times=1000)
def inference(model, example_input):
model(example_input)
inference(model, example_input)
import torch
import torchvision.models as models
from unip.utils.energy import Calculator
calculator = Calculator(cpu=False, device_id=4)
model = models.resnet18()
model.eval().cuda()
def forward_hook_record_input(module, input, output):
setattr(module, "input", input[0])
hook = model.layer1.register_forward_hook(forward_hook_record_input)
example_input = torch.randn(1, 3, 224, 224).cuda()
model(example_input)
hook.remove()
@calculator.measure(times=100)
def inference(model, example_input):
model(example_input)
inference(module, torch.randn_like(module.input))