def pruning_process(model):
print("------------------------- Before pruning --------------------------------")
criterion = nn.CrossEntropyLoss().cuda()
accuracy = util.validate(args, test_loader, model, criterion)
print("------------------------- pruning --------------------------------------")
if args.model_mode == 'c' or args.model_mode =='a':
model.prune_by_percentile( ['conv1'], q=100-58.0)
model.prune_by_percentile( ['conv2'], q=100-22.0)
model.prune_by_percentile( ['conv3'], q=100-34.0)
model.prune_by_percentile( ['conv4'], q=100-36.0)
model.prune_by_percentile( ['conv5'], q=100-53.0)
model.prune_by_percentile( ['conv6'], q=100-24.0)
model.prune_by_percentile( ['conv7'], q=100-42.0)
model.prune_by_percentile( ['conv8'], q=100-32.0)
model.prune_by_percentile( ['conv9'], q=100-27.0)
model.prune_by_percentile( ['conv10'], q=100-34.0)
model.prune_by_percentile( ['conv11'], q=100-35.0)
model.prune_by_percentile( ['conv12'], q=100-29.0)
model.prune_by_percentile( ['conv13'], q=100-36.0)
if args.model_mode == 'd' or args.model_mode == 'a':
model.prune_by_percentile( ['fc1'], q=100-10.0)
model.prune_by_percentile( ['fc2'], q=100-10.0)
model.prune_by_percentile( ['fc3'], q=100-10.0)
print("------------------------- After pruning --------------------------------")
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(args, test_loader, model, criterion)
torch.save(model, f"{args.save_dir}/model_pruned.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_folder}")
util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_pruned.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"accuracy after pruning\t{accuracy}")
util.layer2torch(model, f"{args.save_dir}/{args.out_pruned_folder}")
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}", weight_list)
# Retrain
print("------------------------- Prune and Retrain ----------------------------")
tok="prune_re"
util.initial_train(model, args, train_loader, test_loader, tok, use_cuda=True)
print("------------------------- After Retraining -----------------------------")
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(args, test_loader, model, criterion)
torch.save(model, f"{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_re_folder}")
util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/mmodel_prune_retrain_{args.reepochs}.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"accuracy after prune retrain\t{accuracy}")
util.layer2torch(model, f"{args.save_dir}/{args.out_pruned_re_folder}")
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}", weight_list)
return model
def encoding_process():
util.topic_log("Accuracy before huffman encoding")
util.print_nonzeros(model, args.log_file_path)
top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5))
util.log(args.log_file_path,
f"accuracy before huffman encoding\t{top1_acc} ({top5_acc})")
util.topic_log("encoding")
mesa2_huffman_encode_model(model, args)
def pruning_process():
util.topic_log("Before pruning")
top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5))
util.log(args.log_file_path,
f"before pruning accuracy\t{top1_acc} ({top5_acc})")
util.topic_log("Pruning CNN")
model.prune(args)
util.topic_log("After prune CNN")
util.print_nonzeros(model, args.log_file_path)
top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5))
util.log(args.log_file_path,
f"after pruning accuracy\t{top1_acc} ({top5_acc})")
util.topic_log("Start retrain after prune CNN")
util.train(model, args, train_loader, val_loader, 'prune_retrain')
util.topic_log("After retraining")
top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5))
util.log(args.log_file_path,
f"after pruning and retrain accuracy\t{top1_acc} ({top5_acc})")
def pruning_process(model):
print(
"------------------------- Before pruning --------------------------------"
)
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(val_loader, model, args)
print(
"------------------------- pruning CNN--------------------------------------"
)
model.prune_by_percentile(['conv1'], q=100 - 66.0)
model.prune_by_percentile(['conv2'], q=100 - 12.0)
print(
"------------------------------- After prune CNN ----------------------------"
)
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
prec1 = util.validate(val_loader, model, args)
torch.save(model, f"{args.save_dir}/model_pruned.ptmodel")
util.log(f"{args.save_dir}/{args.log}",
f"weight\t{args.save_dir}/{args.out_pruned_folder}")
util.log(f"{args.save_dir}/{args.log}",
f"model\t{args.save_dir}/model_pruned.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"prune acc\t{prec1}")
util.layer2torch(f"{args.save_dir}/{args.out_pruned_folder}", model)
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}",
weight_list)
print(
"------------------------- start retrain after prune CNN----------------------------"
)
util.initial_train(model, args, train_loader, val_loader, 'prune_re')
print(
"------------------------- After Retraining -----------------------------"
)
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(val_loader, model, args)
torch.save(model,
f"{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel")
util.log(f"{args.save_dir}/{args.log}",
f"weight\t{args.save_dir}/{args.out_pruned_re_folder}")
util.log(
f"{args.save_dir}/{args.log}",
f"model\t{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel")
util.log(f"{args.save_dir}/{args.log}",
f"prune and retrain acc\t{accuracy}")
util.layer2torch(f"{args.save_dir}/{args.out_pruned_re_folder}", model)
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}",
weight_list)
return model
def pruning_process(model):
print("------------------------- Before pruning --------------------------------")
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(val_loader, model, args)
print("------------------------- pruning CNN--------------------------------------")
model.prune_by_percentile( ['conv1'], q=100-58.0)
model.prune_by_percentile( ['conv2'], q=100-22.0)
model.prune_by_percentile( ['conv3'], q=100-34.0)
model.prune_by_percentile( ['conv4'], q=100-36.0)
model.prune_by_percentile( ['conv5'], q=100-53.0)
model.prune_by_percentile( ['conv6'], q=100-24.0)
model.prune_by_percentile( ['conv7'], q=100-42.0)
model.prune_by_percentile( ['conv8'], q=100-32.0)
model.prune_by_percentile( ['conv9'], q=100-27.0)
model.prune_by_percentile( ['conv10'], q=100-34.0)
model.prune_by_percentile( ['conv11'], q=100-35.0)
model.prune_by_percentile( ['conv12'], q=100-29.0)
model.prune_by_percentile( ['conv13'], q=100-36.0)
print("------------------------------- After prune CNN ----------------------------")
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
prec1 = util.validate(val_loader, model, args)
util.save_checkpoint({
'state_dict': model.state_dict(),
'best_prec1': prec1,
}, True, filename=os.path.join(args.save_dir, 'checkpoint_{}_alpha_{}.tar'.format('pruned',args.alpha)))
util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_folder}")
util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_pruned.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"prune acc\t{prec1}")
util.layer2torch(f"{args.save_dir}/{args.out_pruned_folder}" , model)
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}", weight_list)
print("------------------------- start retrain after prune CNN----------------------------")
util.initial_train(model, args, train_loader, val_loader, 'prune_re')
print("------------------------- After Retraining -----------------------------")
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(val_loader, model, args)
util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_re_folder}")
util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel")
util.log(f"{args.save_dir}/{args.log}", f"prune and retrain acc\t{accuracy}")
util.layer2torch(f"{args.save_dir}/{args.out_pruned_re_folder}" , model)
weight_list = util.parameters2list(model.children())
util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}", weight_list)
return model
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(
f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)'
)
return accuracy
# Initial training
print("--- Initial training ---")
train(args.epochs)
accuracy = test()
util.log(args.log, f"initial_accuracy {accuracy}")
torch.save(model, f"saves/initial_model.ptmodel")
print("--- Before pruning ---")
util.print_nonzeros(model)
# Pruning
model.prune_by_std(args.sensitivity)
accuracy = test()
util.log(args.log, f"accuracy_after_pruning {accuracy}")
print("--- After pruning ---")
util.print_nonzeros(model)
# Retrain
print("--- Retraining ---")
optimizer.load_state_dict(initial_optimizer_state_dict) # Reset the optimizer
train(args.epochs)
torch.save(model, f"saves/model_after_retraining.ptmodel")
accuracy = test()
util.log(args.log, f"accuracy_after_retraining {accuracy}")
f = open("result_initial.csv", 'w')
print(epoch_list_initial)
for i in range(0, len(epoch_list_initial)):
f.write(
str(epoch_list_initial[i]) + ',' + str(psnr_list_initial[i]) + "\n")
f.close()
# Pruning
srcnn.prune_by_std(opt.sensitivity)
psnr = test()
util.log(opt.log, "accuracy_after_pruning {}".format(psnr))
print("--- After pruning ---")
util.print_nonzeros(srcnn)
# Retrain
print("--- Retraining ---")
optimizer.load_state_dict(initial_optimizer_state_dict) # Reset the optimizer
epoch_list_retrain = []
psnr_list_retrain = []
for epoch in range(1, opt.epochs + 1):
train(epoch)
psnr = test()
epoch_list_retrain.append(epoch)
psnr_list_retrain.append(psnr)
if (epoch % 100 == 0):
checkpoint(1, epoch)
model = LeNet_300_100(input_size, num_classes)
opt = keras.optimizers.Adam(learning_rate=lr)
model.compile(loss='binary_crossentropy',
optimizer='adam',
metrics=['accuracy'])
#path to save original model
model_path = "lenet300-100.h5"
model, history = train(model, x_train, y_train, batch_size, num_epochs,
model_path) #training
original_loss, original_accuracy = model.evaluate(x_test, y_test,
verbose=0) #testing
print("Result of Original Model:")
print('Original Model loss:', original_loss)
print('Original Model accuracy:', original_accuracy)
print_nonzeros(model)
get_plot(history)
#Iterative training (prune+retrain) in single iteration and we run until we get pruning accuracy >= original model accuracy.
#Threshold is calculated based on standard deviation of weights and sensitivity/quality factor.
for i in range(1, 10):
pretrained_model = keras.models.load_model(model_path)
quality_parameter = 1.6 #sensitivity factor to calculate threshold
#Pruning
for layerid in range(len(pretrained_model.layers)):
layer = pretrained_model.layers[layerid]
weight = layer.get_weights()
if len(weight) > 0:
temp_weight = deepcopy(weight)
def main():
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.gpu is not None:
model = model.cuda(args.gpu)
elif args.distributed:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
else:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
#if args.pretrained and 'alexnet' in args.arch:
# model.load_state_dict(torch.load('saves/new_elt_0.0_0.pth'))
#print(model)
#util.print_model_parameters(model)
# optionally resume from a checkpoint
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 {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=256,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
util.print_nonzeros(model)
validate(val_loader, model, criterion)
return
#if args.pretrained:
# print('Pretrained model evaluation...')
# validate(val_loader, model, criterion)
curves = np.zeros(((args.epochs - args.start_epoch) *
(len(train_loader) // args.print_freq), 4))
valid = np.zeros(((args.epochs - args.start_epoch), 3))
step = 0
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
curves, step = train(train_loader, model, criterion, optimizer, epoch,
args.reg, args.decay, curves, step)
# evaluate on validation set
valid[epoch, 0] = epoch
valid[epoch, 1], valid[epoch, 2] = validate(val_loader, model,
criterion)
# plot training curve
np.savetxt(os.path.join(save_path, 'curves.dat'), curves)
clr1 = (0.5, 0., 0.)
clr2 = (0.0, 0.5, 0.)
fig, ax1 = plt.subplots()
fig2, ax3 = plt.subplots()
ax2 = ax1.twinx()
ax4 = ax3.twinx()
ax1.set_xlabel('steps')
ax1.set_ylabel('Loss', color=clr1)
ax1.tick_params(axis='y', colors=clr1)
ax2.set_ylabel('Reg', color=clr2)
ax2.tick_params(axis='y', colors=clr2)
ax3.set_xlabel('steps')
ax3.set_ylabel('Elt_sparsity', color=clr1)
ax3.tick_params(axis='y', colors=clr1)
start = 0
end = step
markersize = 12
coef = 2.
ax1.plot(curves[start:end, 0],
curves[start:end, 1],
'--',
color=[c * coef for c in clr1],
markersize=markersize)
ax2.plot(curves[start:end, 0],
curves[start:end, 2],
'-',
color=[c * coef for c in clr2],
markersize=markersize)
ax3.plot(curves[start:end, 0],
curves[start:end, 3],
'--',
color=[c * coef for c in clr1],
markersize=markersize)
#ax2.set_ylim(bottom=20, top=100)
ax1.legend(('Train loss'), loc='lower right')
ax2.legend(('Train reg'), loc='lower left')
fig.savefig(os.path.join(save_path, 'loss-vs-steps.pdf'))
#ax4.set_ylim(bottom=20, top=100)
ax3.legend(('Elt_sparsity'), loc='lower right')
fig2.savefig(os.path.join(save_path, 'sparsity-vs-steps.pdf'))
# plot validation curve
np.savetxt(os.path.join(save_path, 'valid.dat'), valid)
fig3, ax5 = plt.subplots()
ax6 = ax5.twinx()
ax5.set_xlabel('epochs')
ax5.set_ylabel('Acc@1', color=clr1)
ax5.tick_params(axis='y', colors=clr1)
ax6.set_ylabel('Acc@5', color=clr2)
ax6.tick_params(axis='y', colors=clr2)
start = 0
end = epoch + 1
markersize = 12
coef = 2.
ax5.plot(valid[start:end, 0],
valid[start:end, 1],
'--',
color=[c * coef for c in clr1],
markersize=markersize)
ax6.plot(valid[start:end, 0],
valid[start:end, 2],
'-',
color=[c * coef for c in clr2],
markersize=markersize)
#ax2.set_ylim(bottom=20, top=100)
ax5.legend(('Acc@1'), loc='lower right')
ax6.legend(('Acc@5'), loc='lower left')
fig3.savefig(os.path.join(save_path, 'accuracy-vs-epochs.pdf'))
torch.save(
model.state_dict(),
os.path.join(
save_path,
'str_' + str(args.decay) + '_' + str(args.reg) + '.pth'))
def main():
global args, best_prec1
args = parser.parse_args()
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
if args.arch.startswith('mobilenetv1'):
model = MobileNetV1()
elif args.arch.startswith('mobilenetv2'):
model = MobileNetV2()
else:
model = models.__dict__[args.arch]()
if args.pretrainedmodeler:
if os.path.isfile(args.pretrainedmodeler):
print("=> loading model '{}'".format(args.pretrainedmodeler))
new_state_dict = OrderedDict()
state_dict = torch.load(args.pretrainedmodeler)
for k, v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print(model)
if not args.distributed:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
else:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
torch.save(model, f"saves/model_after_retraining.ptmodel")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.evaluate:
if os.path.isfile(args.evaluate):
print("=> loading model '{}'".format(args.evaluate))
model.load_state_dict(torch.load(args.evaluate))
else:
print("=> no model found at '{}'".format(args.evaluate))
# optionally resume from a checkpoint
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 {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
num_workers=args.workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer' : optimizer.state_dict(),
}, is_best)
print("--- Before pruning ---")
util.print_nonzeros(model)
def inference(model):
util.print_nonzeros(model, f"{args.save_dir}/{args.log}")
accuracy = util.validate(val_loader, model, args)
print(accuracy)
def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
# Data
print('==> Preparing dataset %s' % args.dataset)
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
if args.dataset == 'cifar10':
dataloader = datasets.CIFAR10
num_classes = 10
else:
dataloader = datasets.CIFAR100
num_classes = 100
trainset = dataloader(root='./data',
train=True,
download=True,
transform=transform_train)
trainloader = data.DataLoader(trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers)
testset = dataloader(root='./data',
train=False,
download=False,
transform=transform_test)
testloader = data.DataLoader(testset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers)
# Model
print("==> creating model '{}'".format(args.arch))
if args.arch.startswith('resnext'):
model = models.__dict__[args.arch](
cardinality=args.cardinality,
num_classes=num_classes,
depth=args.depth,
widen_factor=args.widen_factor,
dropRate=args.drop,
)
elif args.arch.startswith('densenet'):
model = models.__dict__[args.arch](
num_classes=num_classes,
depth=args.depth,
growthRate=args.growthRate,
compressionRate=args.compressionRate,
dropRate=args.drop,
)
elif args.arch.startswith('wrn'):
model = models.__dict__[args.arch](
num_classes=num_classes,
depth=args.depth,
widen_factor=args.widen_factor,
dropRate=args.drop,
)
elif args.arch.endswith('resnet'):
model = models.__dict__[args.arch](
num_classes=num_classes,
depth=args.depth,
block_name=args.block_name,
)
else:
model = models.__dict__[args.arch](num_classes=num_classes)
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Resume
title = 'cifar-10-' + args.arch
model.load_state_dict(
torch.load(os.path.join(args.resume, args.model + '.pth')))
logger = Logger(os.path.join(save_path, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.'
])
print('model loaded')
device = torch.device("cuda")
# Initial training
print("--- Pruning ---")
masks = {}
for name, p in model.named_parameters():
if 'weight' in name:
tensor = p.data.cpu().numpy()
threshold = args.sensitivity
new_mask = np.where(abs(tensor) < threshold, 0., tensor)
mask = np.where(abs(tensor) < threshold, 0., 1.)
masks[name] = torch.from_numpy(mask).float().to(device)
p.data = torch.from_numpy(new_mask).to(device)
util.print_nonzeros(model)
print('Pruned model evaluation...')
test_loss, test_acc = test(testloader, model, criterion, start_epoch,
use_cuda)
print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
if args.evaluate:
return
best_prec1 = test_acc
torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth'))
print("--- Finetuning ---")
# Train and val
for epoch in range(start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
train_loss, train_acc = train(trainloader, model, criterion, optimizer,
epoch, use_cuda, masks)
test_loss, test_acc = test(testloader, model, criterion, epoch,
use_cuda)
# append logger file
logger.append(
[state['lr'], train_loss, test_loss, train_acc, test_acc])
# save model
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
if is_best:
torch.save(model.state_dict(),
os.path.join(save_path, 'finetuned.pth'))
logger.close()
#logger.plot()
#savefig(os.path.join(save_path, 'log.eps'))
print("--- Evaluating ---")
model.load_state_dict(torch.load(os.path.join(save_path, 'finetuned.pth')))
test_loss, test_acc = test(testloader, model, criterion, start_epoch,
use_cuda)
print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
print('Best acc:')
print(best_acc)
def main():
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.gpu is not None:
model = model.cuda(args.gpu)
elif args.distributed:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
else:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
if args.arch.startswith('alexnetv2') and args.pretrained:
pass
else:
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
input_size = 224
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.pretrained is not True:
model.load_state_dict(torch.load(os.path.join(args.resume, args.model+'.pth')))
print('model loaded')
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
num_workers=args.workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
normalize,
])),
batch_size=256, shuffle=False,
num_workers=args.workers, pin_memory=True)
device = torch.device("cuda")
# Initial pruning
print("--- Pruning ---")
masks = {}
for name, p in model.named_parameters():
if 'weight' in name:
tensor = p.data.cpu().numpy()
threshold = args.sensitivity
new_mask = np.where(abs(tensor) < threshold, 0, tensor)
mask = np.where(abs(tensor) < threshold, 0., 1.)
masks[name] = torch.from_numpy(mask).float().to(device)
p.data = torch.from_numpy(new_mask).to(device)
util.print_nonzeros(model)
print('Pruned model evaluation...')
prec1, prec5 = validate(val_loader, model, criterion)
if args.evaluate:
return
best_prec1 = prec5
torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth'))
print("--- Finetuning ---")
curves = np.zeros(((args.epochs-args.start_epoch)*(len(train_loader)//args.print_freq),2))
valid = np.zeros(((args.epochs-args.start_epoch),3))
step = 0
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
curves,step = train(train_loader, model, criterion, optimizer, epoch, curves, step, masks)
# evaluate on validation set
valid[epoch, 0] = epoch
valid[epoch, 1], valid[epoch, 2] = validate(val_loader, model, criterion)
prec5 = valid[epoch, 2]
# plot training curve
np.savetxt(os.path.join(save_path, 'curves.dat'), curves)
clr1 = (0.5, 0., 0.)
clr2 = (0.0, 0.5, 0.)
fig, ax1 = plt.subplots()
ax1.set_xlabel('steps')
ax1.set_ylabel('Loss', color=clr1)
ax1.tick_params(axis='y', colors=clr1)
start = 0
end = step
markersize = 12
coef = 2.
ax1.plot(curves[start:end, 0], curves[start:end, 1], '--', color=[c*coef for c in clr1], markersize=markersize)
#ax2.set_ylim(bottom=20, top=100)
ax1.legend(('Train loss'), loc='lower right')
fig.savefig(os.path.join(save_path, 'loss-vs-steps.pdf'))
# plot validation curve
np.savetxt(os.path.join(save_path, 'valid.dat'), valid)
fig3, ax5 = plt.subplots()
ax6 = ax5.twinx()
ax5.set_xlabel('epochs')
ax5.set_ylabel('Acc@1', color=clr1)
ax5.tick_params(axis='y', colors=clr1)
ax6.set_ylabel('Acc@5', color=clr2)
ax6.tick_params(axis='y', colors=clr2)
start = 0
end = epoch+1
markersize = 12
coef = 2.
ax5.plot(valid[start:end, 0], valid[start:end, 1], '--', color=[c*coef for c in clr1], markersize=markersize)
ax6.plot(valid[start:end, 0], valid[start:end, 2], '-', color=[c*coef for c in clr2], markersize=markersize)
#ax2.set_ylim(bottom=20, top=100)
ax5.legend(('Acc@1'), loc='lower right')
ax6.legend(('Acc@5'), loc='lower left')
fig3.savefig(os.path.join(save_path, 'accuracy-vs-epochs.pdf'))
if prec5 > best_prec1:
torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth'))
best_prec1 = prec5
print('New best performance')
print("--- Evaluating ---")
model.load_state_dict(torch.load(os.path.join(save_path, 'finetuned.pth')))
prec1 = validate(val_loader, model, criterion)
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')
return accuracy
# Initial training
print("--- Initial training ---")
util.log(args.log, "--- Initial training ---")
train(100)
accuracy = test()
util.log(args.log, f"initial_accuracy {accuracy}")
torch.save(model, f"saves/initial_model.ptmodel")
print("--- Before pruning ---")
util.log(args.log, "--- Before pruning ---")
util.log(args.log, util.print_nonzeros(model))
original_accuracy = accuracy
loop_cnt = 0
minus_cnt = 0
start = time.time()
while (time.time() - start) / 3600 < 2:
print("Iteration", loop_cnt)
util.log(args.log, f"Iteration {loop_cnt}")
loop_cnt += 1
# Pruning
model.prune_by_std(args.sensitivity)
accuracy = test()
util.log(args.log, f"accuracy_after_pruning {accuracy}")
print("--- After pruning ---")
