nni模型剪枝
本文参考:
【模型部署】NNI：剪枝和量化_Jackilina_Stone的博客-CSDN博客_nni剪枝
模型剪枝入门 — Neural Network Intelligence
NNI：Neural Network Intelligence，是一个轻量但强大的自动机器学习工具包，能帮助用户自动进行特征工程、神经网络架构搜索、超参调优以及模型压缩。
一、安装NNI
pip install nni安装完成后检查nni版本：
nnictl --version二、剪枝操作
之前已经写过两篇关于模型剪枝的方法，分别是：
模型剪枝初级方法_benben044的博客-CSDN博客
Slimming剪枝方法_benben044的博客-CSDN博客
VGG模型代码：
import torch
import torch.nn as nn
import math
from torch.autograd import Variable
class vgg(nn.Module):
def __init__(self, dataset='cifar10', init_weights=True, cfg=None):
super(vgg, self).__init__()
if cfg is None:
cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512]
self.feature = self.make_layers(cfg, True)
if dataset == 'cifar10':
num_classes = 10
elif dataset == 'cifar100':
num_classes = 100
self.classifier = nn.Linear(cfg[-1], num_classes)
if init_weights:
self._initialize_weights()
def make_layers(self, cfg, batch_norm=False):
layers = []
in_channels = 3
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1, bias=False)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
def forward(self, x):
x = self.feature(x)
x = nn.AvgPool2d(2)(x)
x = x.view(x.size(0), -1)
y = self.classifier(x)
return y
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(0.5)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
if __name__ == '__main__':
net = vgg()
x = Variable(torch.FloatTensor(16, 3, 40, 40))
y = net(x)
print(y.data.shape)通过nni模型剪枝的代码：
import torch
import torch.optim as optim
from vgg import vgg
from tqdm import tqdm
from torchvision import datasets, transforms
from torch.autograd import Variable
import torch.nn.functional as F
from nni.compression.pytorch.pruning import L1NormPruner
from nni.compression.pytorch.speedup import ModelSpeedup
import shutil
config_list = [{
'sparsity_per_layer': 0.5,
'op_types': ['Conv2d']
}]
learning_rate = 0.1
momentum = 0.9
weight_decay = 1e-4
epochs = 6
log_interval = 100
batch_size = 100
sparsity_regularization = True
scale_sparse_rate = 0.0001
prune_checkpoint_path = 'nni_pruned_checkpoint.pth.tar'
prune_best_model_path = 'nni_pruned_model_best.pth.tar'
model = vgg()
model.cuda()
print("--------------raw model--------------")
print(model)
pruner = L1NormPruner(model, config_list)
print("--------------pruned model--------------")
print(model)
# compress the model and generate the masks
_, masks = pruner.compress()
# show the masks sparsity
for name, mask in masks.items():
print(name, ' sparsity: ', '{:.2f}'.format(mask['weight'].sum() / mask['weight'].numel()))
# need to unwarp the model, if the model is wrawpped before speedup
pruner._unwrap_model()
# speedup the model
ModelSpeedup(model, torch.rand(16, 3, 40, 40).cuda(), masks).speedup_model()
print("--------------after sppedup--------------")
print(model)
# retrain
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
train_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('D:\\ai_data\\cifar10', train=True, download=True,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])),
batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('D:\\ai_data\\cifar10', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])),
batch_size=batch_size, shuffle=True)
def train(epoch):
model.train()
for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
data, target = data.cuda(), target.cuda()
data, target = Variable(data), Variable(target)
optimizer.zero_grad()
output = model(data)
loss = F.cross_entropy(output, target)
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.1f}%)]\tLoss: {:.6f}'.format(epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))
def test():
model.eval()
test_loss = 0
correct = 0
for data, target in tqdm(test_loader):
data , target = data.cuda(), target.cuda()
data, target = Variable(data), Variable(target)
output = model(data)
test_loss += F.cross_entropy(output, target, size_average=False).item()
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(target.data.view_as(pred)).cpu().sum()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
return correct / float(len(test_loader.dataset))
def save_checkpoint(state, is_best, filename=prune_checkpoint_path):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename, prune_best_model_path)
best_prec = 0
for epoch in range(epochs):
train(epoch)
prec = test()
is_best = prec > best_prec
best_prec = max(prec, best_prec)
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec': best_prec,
'optimizer': optimizer.state_dict()
}, is_best)
在运行相同的迭代次数情况下，nni的准确率没有slimming高