Initial

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Mingjie Sun
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,5 @@
+# Rethinking the Value of Network Pruning (Pytorch)
+This repository contains a pytorch implementation of the paper Rethinking the Value of Network Pruning.
+
+## Contact
+sunmj15 at gmail.com liuzhuangthu at gmail.com

cifar/README.md

@@ -0,0 +1,7 @@
+# CIFAR Experiments
+This directory contains all the CIFAR experiments in the paper, where there are three pruning methods in total:  
+1.  L1-norm based channel pruning  
+2.  Network Slimming  
+3.  Non-structured weight level pruning 
+
+For each method, we give example commands for baseline training, finetuning, scratch-E training and scratch-B training.

cifar/l1-norm-pruning/README.md

@@ -0,0 +1,48 @@
+# Pruning Filters For Efficient ConvNets
+
+This directory contains a pytorch re-implementation of all CIFAR experiments of the following paper  
+[Pruning Filters for Efficient ConvNets](https://arxiv.org/abs/1608.08710) (ICLR 2017).
+
+## Dependencies
+torch v0.3.1, torchvision v0.2.0
+
+## Baseline 
+
+The `dataset` argument specifies which dataset to use: `cifar10` or `cifar100`. The `arch` argument specifies the architecture to use: `vgg` or `resnet`. The depth is chosen to be the same as the networks used in the paper.
+```shell
+python main.py --dataset cifar10 --arch vgg --depth 16
+python main.py --dataset cifar10 --arch resnet --depth 56
+python main.py --dataset cifar10 --arch resnet --depth 110
+```
+
+## Prune
+
+```shell
+python vggprune.py --dataset cifar10 --model [PATH TO THE MODEL] --save [DIRECTORY TO STORE RESULT]
+python res56prune.py --dataset cifar10 -v A --model [PATH TO THE MODEL] --save [DIRECTORY TO STORE RESULT]
+python res110prune.py --dataset cifar10 -v A --model [PATH TO THE MODEL] --save [DIRECTORY TO STORE RESULT]
+```
+Here in `res56prune.py` and `res110prune.py`, the `-v` argument is `A` or `B`, which refers to the naming of the pruned model in the original paper. The pruned model will be named `pruned.pth.tar`.
+
+## Fine-tune
+
+```shell
+python main_finetune.py --refine [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch vgg --depth 16 
+python main_finetune.py --refine [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 56 
+python main_finetune.py --refine [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 110 
+```
+
+## Scratch-E
+```
+python main_E.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch vgg --depth 16
+python main_E.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 56 
+python main_E.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 110  
+```
+
+## Scratch-B
+```
+python main_B.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch vgg --depth 16
+python main_B.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 56
+python main_B.py --scratch [PATH TO THE PRUNED MODEL] --dataset cifar10 --arch resnet --depth 110
+```
+

cifar/l1-norm-pruning/compute_flops.py

@@ -0,0 +1,112 @@
+import numpy as np
+
+import torch
+import torchvision
+import torch.nn as nn
+from torch.autograd import Variable
+
+
+def print_model_param_nums(model=None, multiply_adds=True):
+    if model == None:
+        model = torchvision.models.alexnet()
+    total = sum([param.nelement() for param in model.parameters()])
+    print('  + Number of params: %.2fM' % (total / 1e6))
+
+def print_model_param_flops(model=None, input_res=224, multiply_adds=True):
+
+    prods = {}
+    def save_hook(name):
+        def hook_per(self, input, output):
+            prods[name] = np.prod(input[0].shape)
+        return hook_per
+
+    list_1=[]
+    def simple_hook(self, input, output):
+        list_1.append(np.prod(input[0].shape))
+    list_2={}
+    def simple_hook2(self, input, output):
+        list_2['names'] = np.prod(input[0].shape)
+
+    list_conv=[]
+    def conv_hook(self, input, output):
+        batch_size, input_channels, input_height, input_width = input[0].size()
+        output_channels, output_height, output_width = output[0].size()
+
+        kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups)
+        bias_ops = 1 if self.bias is not None else 0
+
+        params = output_channels * (kernel_ops + bias_ops)
+        flops = (kernel_ops * (2 if multiply_adds else 1) + bias_ops) * output_channels * output_height * output_width * batch_size
+
+        list_conv.append(flops)
+
+    list_linear=[]
+    def linear_hook(self, input, output):
+        batch_size = input[0].size(0) if input[0].dim() == 2 else 1
+
+        weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
+        bias_ops = self.bias.nelement()
+
+        flops = batch_size * (weight_ops + bias_ops)
+        list_linear.append(flops)
+
+    list_bn=[]
+    def bn_hook(self, input, output):
+        list_bn.append(input[0].nelement() * 2)
+
+    list_relu=[]
+    def relu_hook(self, input, output):
+        list_relu.append(input[0].nelement())
+
+    list_pooling=[]
+    def pooling_hook(self, input, output):
+        batch_size, input_channels, input_height, input_width = input[0].size()
+        output_channels, output_height, output_width = output[0].size()
+
+        kernel_ops = self.kernel_size * self.kernel_size
+        bias_ops = 0
+        params = 0
+        flops = (kernel_ops + bias_ops) * output_channels * output_height * output_width * batch_size
+
+        list_pooling.append(flops)
+
+    list_upsample=[]
+    # For bilinear upsample
+    def upsample_hook(self, input, output):
+        batch_size, input_channels, input_height, input_width = input[0].size()
+        output_channels, output_height, output_width = output[0].size()
+
+        flops = output_height * output_width * output_channels * batch_size * 12
+        list_upsample.append(flops)
+
+    def foo(net):
+        childrens = list(net.children())
+        if not childrens:
+            if isinstance(net, torch.nn.Conv2d):
+                net.register_forward_hook(conv_hook)
+            if isinstance(net, torch.nn.Linear):
+                net.register_forward_hook(linear_hook)
+            if isinstance(net, torch.nn.BatchNorm2d):
+                net.register_forward_hook(bn_hook)
+            if isinstance(net, torch.nn.ReLU):
+                net.register_forward_hook(relu_hook)
+            if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
+                net.register_forward_hook(pooling_hook)
+            if isinstance(net, torch.nn.Upsample):
+                net.register_forward_hook(upsample_hook)
+            return
+        for c in childrens:
+            foo(c)
+
+    if model == None:
+        model = torchvision.models.alexnet()
+    foo(model)
+    input = Variable(torch.rand(3, 3, input_res, input_res), requires_grad = True)
+    out = model(input)
+
+
+    total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling) + sum(list_upsample))
+
+    print('  + Number of FLOPs: %.5fG' % (total_flops / 3 / 1e9))
+
+    return total_flops

cifar/l1-norm-pruning/main.py

@@ -0,0 +1,176 @@
+from __future__ import print_function
+import argparse
+import numpy as np
+import os
+import shutil
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.autograd import Variable
+
+import models
+
+
+# Training settings
+parser = argparse.ArgumentParser(description='PyTorch Slimming CIFAR training')
+parser.add_argument('--dataset', type=str, default='cifar100',
+                    help='training dataset (default: cifar100)')
+parser.add_argument('--batch-size', type=int, default=64, metavar='N',
+                    help='input batch size for training (default: 64)')
+parser.add_argument('--test-batch-size', type=int, default=256, metavar='N',
+                    help='input batch size for testing (default: 256)')
+parser.add_argument('--epochs', type=int, default=160, metavar='N',
+                    help='number of epochs to train (default: 160)')
+parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                    help='manual epoch number (useful on restarts)')
+parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
+                    help='learning rate (default: 0.1)')
+parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+                    help='SGD momentum (default: 0.9)')
+parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
+                    metavar='W', help='weight decay (default: 1e-4)')
+parser.add_argument('--resume', default='', type=str, metavar='PATH',
+                    help='path to latest checkpoint (default: none)')
+parser.add_argument('--no-cuda', action='store_true', default=False,
+                    help='disables CUDA training')
+parser.add_argument('--seed', type=int, default=1, metavar='S',
+                    help='random seed (default: 1)')
+parser.add_argument('--log-interval', type=int, default=100, metavar='N',
+                    help='how many batches to wait before logging training status')
+parser.add_argument('--save', default='./logs', type=str, metavar='PATH',
+                    help='path to save prune model (default: current directory)')
+parser.add_argument('--arch', default='vgg', type=str, 
+                    help='architecture to use')
+parser.add_argument('--depth', default=16, type=int,
+                    help='depth of the neural network')
+
+args = parser.parse_args()
+args.cuda = not args.no_cuda and torch.cuda.is_available()
+
+torch.manual_seed(args.seed)
+if args.cuda:
+    torch.cuda.manual_seed(args.seed)
+
+if not os.path.exists(args.save):
+    os.makedirs(args.save)
+
+kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
+if args.dataset == 'cifar10':
+    train_loader = torch.utils.data.DataLoader(
+        datasets.CIFAR10('./data.cifar10', train=True, download=True,
+                       transform=transforms.Compose([
+                           transforms.Pad(4),
+                           transforms.RandomCrop(32),
+                           transforms.RandomHorizontalFlip(),
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+                       ])),
+        batch_size=args.batch_size, shuffle=True, **kwargs)
+    test_loader = torch.utils.data.DataLoader(
+        datasets.CIFAR10('./data.cifar10', train=False, transform=transforms.Compose([
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+                       ])),
+        batch_size=args.test_batch_size, shuffle=True, **kwargs)
+else:
+    train_loader = torch.utils.data.DataLoader(
+        datasets.CIFAR100('./data.cifar100', train=True, download=True,
+                       transform=transforms.Compose([
+                           transforms.Pad(4),
+                           transforms.RandomCrop(32),
+                           transforms.RandomHorizontalFlip(),
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+                       ])),
+        batch_size=args.batch_size, shuffle=True, **kwargs)
+    test_loader = torch.utils.data.DataLoader(
+        datasets.CIFAR100('./data.cifar100', train=False, transform=transforms.Compose([
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+                       ])),
+        batch_size=args.test_batch_size, shuffle=True, **kwargs)
+
+model = models.__dict__[args.arch](dataset=args.dataset, depth=args.depth)
+
+if args.cuda:
+    model.cuda()
+
+optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+
+if args.resume:
+    if os.path.isfile(args.resume):
+        print("=> loading checkpoint '{}'".format(args.resume))
+        checkpoint = torch.load(args.resume)
+        args.start_epoch = checkpoint['epoch']
+        best_prec1 = checkpoint['best_prec1']
+        model.load_state_dict(checkpoint['state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        print("=> loaded checkpoint '{}' (epoch {}) Prec1: {:f}"
+              .format(args.resume, checkpoint['epoch'], best_prec1))
+    else:
+        print("=> no checkpoint found at '{}'".format(args.resume))
+
+def train(epoch):
+    model.train()
+    avg_loss = 0.
+    train_acc = 0.
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if args.cuda:
+            data, target = data.cuda(), target.cuda()
+        data, target = Variable(data), Variable(target)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.cross_entropy(output, target)
+        avg_loss += loss.data[0]
+        pred = output.data.max(1, keepdim=True)[1]
+        train_acc += pred.eq(target.data.view_as(pred)).cpu().sum()
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.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.data[0]))
+
+def test():
+    model.eval()
+    test_loss = 0
+    correct = 0
+    for data, target in test_loader:
+        if args.cuda:
+            data, target = data.cuda(), target.cuda()
+        data, target = Variable(data, volatile=True), Variable(target)
+        output = model(data)
+        test_loss += F.cross_entropy(output, target, size_average=False).data[0] # sum up batch loss
+        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
+        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, filepath):
+    torch.save(state, os.path.join(filepath, 'checkpoint.pth.tar'))
+    if is_best:
+        shutil.copyfile(os.path.join(filepath, 'checkpoint.pth.tar'), os.path.join(filepath, 'model_best.pth.tar'))
+
+best_prec1 = 0.
+for epoch in range(args.start_epoch, args.epochs):
+    if epoch in [args.epochs*0.5, args.epochs*0.75]:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] *= 0.1
+    train(epoch)
+    prec1 = test()
+    is_best = prec1 > best_prec1
+    best_prec1 = max(prec1, best_prec1)
+    save_checkpoint({
+        'epoch': epoch + 1,
+        'state_dict': model.state_dict(),
+        'best_prec1': best_prec1,
+        'optimizer': optimizer.state_dict(),
+        'cfg': model.cfg
+    }, is_best, filepath=args.save)