import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from collections import OrderedDict


# 演示随机非结构化剪枝
def dm01():
    linear = nn.Linear(2, 3)
    print("linear-->", linear.weight)
    # model = prune.random_unstructured(linear, 'weight', amount=2)
    model = prune.l1_unstructured(linear, 'weight', amount=2)
    print("model-->", model.weight)


# 演示全局非结构化剪枝
def dm02():
    net = nn.Sequential(OrderedDict([
        ('first', nn.Linear(3, 4)),
        ('second', nn.Linear(4, 2)),
    ]))
    print("net1-->", net)
    for model in net:
        print("model-->", model.weight)
    parameters_to_prune = ((net.first, 'weight'),
                           (net.second, 'weight'))
    # parameters_to_prune：待剪枝的参数
    # pruning_method：剪枝的方式，L1Unstructured表示非结构化剪枝（常用）
    # amount：如果是小数，则表示比例，如果是整数，则表示数量
    prune.global_unstructured(parameters_to_prune,
                              pruning_method=prune.L1Unstructured,
                              amount = 0.2)
    print("net2-->", net)
    for model in net:
        print("model-->", model.weight)


if __name__ == '__main__':
    # dm01()
    dm02()