import torch
import sys
import numpy as np
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data.sampler import SubsetRandomSampler
from torchvision import transforms
from model.resnet_cifar10 import BasicBlock
from pruner.fp_mbnetv2 import FilterPrunerMBNetV2
from pruner.fp_resnet import FilterPrunerResNet
import argparse
def measure_model(model, pruner, img_size):
pruner.reset()
model.eval()
pruner.forward(torch.zeros((1,3,img_size,img_size), device='cuda'))
cur_flops = pruner.cur_flops
cur_size = pruner.cur_size
return cur_flops, cur_size
def save_checkpoint(state, is_best, filename='checkpoint'):
if is_best:
torch.save(state, '{}_best.pth.tar'.format(filename))
def get_valid_flops(model, cbns, out_maps):
lastConv = None
residual_chain = {}
chain_max_dim = 0
for m in model.modules():
if isinstance(m, BasicBlock):
residual_chain[lastConv] = m.conv[3]
lastConv = m.conv[3]
chain_max_dim = np.maximum(chain_max_dim, lastConv.weight.size(1))
if isinstance(m, nn.Conv2d):
lastConv = m
chain_max_dim = lastConv.weight.size(1)
# Deal with the chain first
mask = np.zeros(chain_max_dim)
for key in residual_chain:
conv = residual_chain[key]
target_idx = cbns[0].index(conv)
target_bn = cbns[1][target_idx]
cur_mask = target_bn.weight.data.cpu().numpy()
cur_mask = np.concatenate((cur_mask, np.zeros(chain_max_dim - len(cur_mask))))
mask = np.logical_or(mask, cur_mask)
flops = 0
for idx, (conv, bn) in enumerate(zip(*cbns)):
if conv in residual_chain:
cur_mask = mask[:bn.weight.size(0)]
valid_output = np.sum(cur_mask)
if idx == 0:
valid_input = conv.weight.size(1)
else:
valid_input = (torch.abs(cbns[1][idx-1].weight) > 0).sum().item()
else:
valid_output = (torch.abs(bn.weight) > 0).sum().item()
cur_mask = mask[:cbns[1][idx-1].weight.size(0)]
valid_input = np.sum(cur_mask)
flops += out_maps[idx][0] * out_maps[idx][1] * valid_output * valid_input * conv.weight.size(2) * conv.weight.size(3) / conv.groups
return flops
def get_cbns(model):
convs = []
bns = []
for m in model.modules():
# store the information for batchnorm
if isinstance(m, nn.Conv2d):
convs.append(m)
elif isinstance(m, nn.BatchNorm2d):
bns.append(m)
return convs, bns
def regularizer(model, constraint='size', cbns=None, maps=None):
# build kv map
if cbns is None:
cbns = get_cbns(model)
else:
G = torch.zeros([1], requires_grad=True).cuda()
for idx, (conv, bn) in enumerate(zip(*cbns)):
if idx < len(cbns[0])-1:
gamma_prev = torch.abs(bn.weight)
A = (gamma_prev > 0)
gamma_now = torch.abs(cbns[1][idx+1].weight)
B = (gamma_now > 0)
if constraint == 'size':
cost = cbns[0][idx+1].weight.size(2)*cbns[0][idx+1].weight.size(3)
elif constraint == 'flops':
assert maps is not None, 'Output Map is None!'
cost = 2 * maps[idx+1][0] * maps[idx+1][0] * cbns[0][idx+1].weight.size(2) * cbns[0][idx+1].weight.size(3)
G = G + cost * (gamma_prev.sum()*B.sum().type_as(gamma_prev) + gamma_now.sum()*A.sum().type_as(gamma_now))
return G
def num_alive_filters(model):
cnt = 0
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
cnt = cnt + (torch.abs(m.weight) > 0).sum().item()
return cnt
# Truncate small beta and enforce depth-wise in-out numbers
def truncate_smallbeta(model, cbns):
lastConv = None
residual_chain = {}
chain_max_dim = 0
for m in model.modules():
if isinstance(m, BasicBlock):
residual_chain[lastConv] = m.conv[3]
lastConv = m.conv[3]
chain_max_dim = np.maximum(chain_max_dim, lastConv.weight.size(1))
if isinstance(m, nn.Conv2d):
lastConv = m
chain_max_dim = lastConv.weight.size(1)
# Deal with the chain first
mask = np.zeros(chain_max_dim)
for key in residual_chain:
conv = residual_chain[key]
target_idx = cbns[0].index(conv)
target_bn = cbns[1][target_idx]
cur_mask = target_bn.weight.data.cpu().numpy()
zero_idx = np.abs(cur_mask) < 0.01
cur_mask[zero_idx] = 0
cur_mask = np.concatenate((cur_mask, np.zeros(chain_max_dim - len(cur_mask))))
mask = np.logical_or(mask, cur_mask)
for idx, (conv, bn) in enumerate(zip(*cbns)):
weights = bn.weight.data.cpu().numpy()
bias = bn.bias.data.cpu().numpy()
if conv in residual_chain:
cur_mask = mask[:weights.shape[0]]
weights *= cur_mask
bias *= cur_mask
else:
idx_out = np.abs(weights) < 0.01
weights[idx_out] = 0
bias[idx_out] = 0
bn.weight.data = torch.from_numpy(weights).cuda()
bn.bias.data = torch.from_numpy(bias).cuda()
def test(model, loader):
model.eval()
total = 0
top1 = 0
total_loss = 0
criterion = torch.nn.CrossEntropyLoss()
for i, (batch, label) in enumerate(loader):
batch, label = batch.to('cuda'), label.to('cuda')
total += batch.size(0)
out = model(batch)
total_loss += criterion(out, label).item()
_, pred = out.max(dim=1)
top1 += pred.eq(label).sum()
return float(top1)/total*100, total_loss/total
def train_epoch(model, optim, criterion, loader, lbda=None, cbns=None, maps=None, constraint=None):
model.train()
total = 0
top1 = 0
for i, (batch, label) in enumerate(loader):
optim.zero_grad()
batch, label = batch.to('cuda'), label.to('cuda')
total += batch.size(0)
out = model(batch)
_, pred = out.max(dim=1)
top1 += pred.eq(label).sum()
if constraint:
reg = lbda * regularizer(model, constraint, cbns, maps)
loss = criterion(out, label) + reg
else:
loss = criterion(out, label)
loss.backward()
optim.step()
if (i % 100 == 0) or (i == len(loader)-1):
print('Train | Batch ({}/{}) | Top-1: {:.2f} ({}/{})'.format(
i+1, len(loader),
float(top1)/total*100, top1, total))
if constraint:
truncate_smallbeta(model, cbns)
def train(model, train_loader, val_loader, epochs=10, lr=1e-2, name=''):
model = model.to('cuda')
model.train()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4, nesterov=True)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [int(epochs*0.3), int(epochs*0.6), int(epochs*0.8)], gamma=0.2)
criterion = torch.nn.CrossEntropyLoss()
for e in range(epochs):
train_epoch(model, optimizer, criterion, train_loader)
top1, val_loss = test(model, val_loader)
print('Epoch {} | Top-1: {:.2f}'.format(e, top1))
torch.save(model, 'ckpt/{}_best.t7'.format(name))
scheduler.step()
return model
def train_mask(model, train_loader, val_loader, pruner, epochs=10, lr=1e-2, lbda=1.3*1e-8, cbns=None, maps=None, constraint='flops'):
model = model.to('cuda')
model.train()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4, nesterov=True)
criterion = torch.nn.CrossEntropyLoss()
for e in range(epochs):
print('Epoch {}'.format(e))
train_epoch(model, optimizer, criterion, train_loader, lbda, cbns, maps, constraint)
top1, _ = test(model, val_loader)
print('#Filters: {}, #FLOPs: {:.2f}M | Top-1: {:.2f}'.format(num_alive_filters(model), pruner.get_valid_flops()/1000000., top1))
return model
def prune_model(model, cbns, pruner):
filters_to_prune_per_layer = pruner.get_valid_filters()
prune_targets = pruner.pack_pruning_target(filters_to_prune_per_layer, get_segment=True, progressive=True)
layers_prunned = {}
for layer_index, filter_index in prune_targets:
if layer_index not in layers_prunned:
layers_prunned[layer_index] = 0
layers_prunned[layer_index] = layers_prunned[layer_index] + (filter_index[1]-filter_index[0]+1)
print('Layers that will be prunned: {}'.format(sorted(layers_prunned.items())))
print('Prunning filters..')
for layer_index, filter_index in prune_targets:
pruner.prune_conv_layer_segment(layer_index, filter_index)
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--datapath", type=str, default='/data')
parser.add_argument("--dataset", type=str, default='torchvision.datasets.CIFAR10')
parser.add_argument("--epoch", type=int, default=60)
parser.add_argument("--name", type=str, default='ft_mbnetv2')
parser.add_argument("--model", type=str, default='ft_mbnetv2')
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--lr", type=float, default=1e-2)
parser.add_argument("--lbda", type=float, default=3e-9)
parser.add_argument("--prune_away", type=float, default=0.5, help='The constraint level in portion to the original network, e.g. 0.5 is prune away 50%')
parser.add_argument("--constraint", type=str, default='flops')
parser.add_argument("--large_input", action='store_true', default=False)
parser.add_argument("--no_grow", action='store_true', default=False)
parser.add_argument("--pruner", type=str, default='FilterPrunnerResNet', help='Different network require differnt pruner implementation')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
print(args)
model = torch.load(args.model)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = eval(args.dataset)(args.datapath, True, transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = eval(args.dataset)(args.datapath, True, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
num_train = len(train_set)
indices = list(range(num_train))
split = int(np.floor(0.1 * num_train))
np.random.seed(98)
np.random.shuffle(indices)
train_idx, valid_idx = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(valid_idx)
test_set = eval(args.dataset)(args.datapath, False, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=args.batch_size, shuffle=True,
num_workers=0, pin_memory=True
)
val_loader = torch.utils.data.DataLoader(
val_set, batch_size=args.batch_size, sampler=valid_sampler,
num_workers=0, pin_memory=True
)
test_loader = torch.utils.data.DataLoader(
test_set, batch_size=125, shuffle=False,
num_workers=0, pin_memory=False
)
if 'CIFAR10' in args.dataset:
train_set.num_classes = 10
elif 'CIFAR100' in args.dataset:
train_set.num_classes = 100
pruner = eval(args.pruner)(model, 'l2_weight', num_cls=train_set.num_classes)
flops, num_params = measure_model(pruner.model, pruner, 32)
maps = pruner.omap_size
cbns = get_cbns(pruner.model)
print('Before Pruning | FLOPs: {:.3f}M | #Params: {:.3f}M'.format(flops/1000000., num_params/1000000.))
train_mask(pruner.model, train_loader, val_loader, pruner, epochs=args.epoch, lr=1e-3, lbda=args.lbda, cbns=cbns, maps=maps, constraint=args.constraint)
target = int((1.-args.prune_away)*flops)
print('Target ({}): {:.3f}M'.format(args.constraint, target/1000000.))
prune_model(pruner.model, cbns, pruner)
flops, num_params = measure_model(pruner.model, pruner, 32)
print('After Pruning | FLOPs: {:.3f}M | #Params: {:.3f}M'.format(flops/1000000., num_params/1000000.))
if args.no_grow:
train(model, train_loader, test_loader, epochs=args.epoch, lr=args.lr, name='{}_pregrow'.format(args.name))
else:
if flops < target:
ratio = pruner.get_uniform_ratio(target)
print(ratio)
pruner.uniform_grow(ratio)
flops, num_params = measure_model(pruner.model, pruner, 32)
print('After Growth | FLOPs: {:.3f}M | #Params: {:.3f}M'.format(flops/1000000., num_params/1000000.))
train(pruner.model, train_loader, test_loader, epochs=args.epoch, lr=args.lr, name=args.name)
else:
print('Over constraint ({:.3f}M > {:.3f}M), no growth'.format(flops/1000000., target/1000000.))
