基于pytorch的模型压缩和模型剪枝Model Prune示例

神经网络和卷积神经网络的模型剪枝Model Prune

1,神经网络和卷积神经网络模型剪枝方法。

2,可指定剪枝率进行定向剪枝,并输出剪枝后参数统计和finetune。

3,支持MLP, Lenet, Alexnet, VGG, GoogleNet系列, Resnet系列,MobileNet系列的剪枝。

下载地址:下载地址

#!/usr/bin/env python
# coding: utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
import torch.utils.data
import numpy as np
import math
def to_var(x, requires_grad=False):
    """
    Automatically choose cpu or cuda
    """
    if torch.cuda.is_available():
        x = x.cuda()
    return x.clone().detach().requires_grad_(requires_grad)

class MaskedConv2d(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=True):
        super(MaskedConv2

你可能感兴趣的:(pytorch量化感知训练,稀疏训练,模型剪枝学习教程,人工智能,python,神经网络)