def main():
# get options
opt = BaseOptions().parse()
# basic settings
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = 'cuda'
torch.backends.cudnn.benchmark = True
else:
device = 'cpu'
##################### Get Dataloader ####################
_, dataloader_test = custom_get_dataloaders(opt)
# dummy_input is sample input of dataloaders
if hasattr(dataloader_test, 'dataset'):
dummy_input = dataloader_test.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Evaluate Baseline Model ####################
net = ModelWrapper(opt)
net = net.to(device)
net.parallel(opt.gpu_ids)
flops_before, params_before = model_summary(net.get_compress_part(),
dummy_input)
del net
##################### Evaluate Pruned Model ####################
net = ModelWrapper(opt)
net.load_checkpoint(opt.pruned_model)
net = net.to(device)
flops_after, params_after = model_summary(net.get_compress_part(),
dummy_input)
net.parallel(opt.gpu_ids)
acc_after = net.get_eval_scores(dataloader_test)
#################### Report #####################
print('######### Report #########')
print('Model:{}'.format(opt.model_name))
print('Checkpoint:{}'.format(opt.pruned_model))
print('FLOPs of Original Model:{:.3f}G;Params of Original Model:{:.2f}M'.
format(flops_before / 1e9, params_before / 1e6))
print('FLOPs of Pruned Model:{:.3f}G;Params of Pruned Model:{:.2f}M'.
format(flops_after / 1e9, params_after / 1e6))
print('Top-1 Acc of Pruned Model on {}:{}'.format(opt.dataset_name,
acc_after['accuracy']))
print('##########################')
def get_eval(opt,net_path,args):
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
load(net,net_path)
##################### Pruning Strategy Generation ###############
compression_scheduler = distiller.file_config(
net.get_compress_part(), net.optimizer, opt.compress_schedule_path
)
compression_scheduler = SetCompressionScheduler(
compression_scheduler, args["channel_config"]
)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=args["dummy_input"])
flops_after, params_after = model_summary(net.get_compress_part(), args["dummy_input"])
ratio = flops_after / args["flops_before"]
net = net.to(args["device"])
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
with torch.no_grad():
for index, sample in enumerate(args["dataloader_train"]):
_ = net.get_loss(sample)
if index > 100:
break
strategy_score = net.get_eval_scores(args["dataloader_val"])["accuracy"]
del net
return 1-strategy_score,ratio
def main(opt):
# basic settings
os.environ["CUDA_VISIBLE_DEVICES"]=str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = 'cuda'
torch.backends.cudnn.benchmark = True
else:
device = 'cpu'
##################### Get Dataloader ####################
dataloader_train, dataloader_val = custom_get_dataloaders(opt)
# dummy_input is sample input of dataloaders
if hasattr(dataloader_val, 'dataset'):
dummy_input = dataloader_val.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
net.load_checkpoint(opt.checkpoint)
flops_before, params_before = model_summary(net.get_compress_part(), dummy_input)
##################### Pruning Strategy Generation ###############
compression_scheduler = distiller.file_config(net.get_compress_part(), net.optimizer, opt.compress_schedule_path)
num_layer = len(compression_scheduler.policies[1])
channel_config = get_pruning_strategy(opt, num_layer) # pruning strategy
compression_scheduler = random_compression_scheduler(compression_scheduler, channel_config)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=dummy_input)
print(net)
flops_after, params_after = model_summary(net.get_compress_part(), dummy_input)
ratio = flops_after / flops_before
print('FLOPs ratio:', ratio)
if ratio < opt.flops_target - 0.005 or ratio > opt.flops_target + 0.005:
# illegal pruning strategy
return
net = net.to(device)
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
with torch.no_grad():
for index, sample in enumerate(tqdm(dataloader_train, leave=False)):
_ = net.get_loss(sample)
if index > 50:
break
strategy_score = net.get_eval_scores(dataloader_val)['accuracy']
#################### Save Pruning Strategy and Score #########
log_file = open(opt.output_file, 'a+')
log_file.write("{} {} ".format(strategy_score, ratio))
for item in channel_config:
log_file.write("{} ".format(str(item)))
log_file.write('\n')
log_file.close()
print('Eval Score:{}'.format(strategy_score))
def main(opt):
# basic settings
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = "cuda"
torch.backends.cudnn.benchmark = True
else:
device = "cpu"
##################### Get Dataloader ####################
dataloader_train, dataloader_val = custom_get_dataloaders(opt)
# dummy_input is sample input of dataloaders
if hasattr(dataloader_val, "dataset"):
dummy_input = dataloader_val.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
net.load_checkpoint(opt.checkpoint)
flops_before, params_before = model_summary(net.get_compress_part(),
dummy_input)
##################### Load Pruning Strategy ###############
compression_scheduler = distiller.file_config(net.get_compress_part(),
net.optimizer,
opt.compress_schedule_path)
channel_config = get_channel_config(opt.search_result,
opt.strategy_id) # pruning strategy
compression_scheduler = random_compression_scheduler(
compression_scheduler, channel_config)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=dummy_input)
flops_after, params_after = model_summary(net.get_compress_part(),
dummy_input)
ratio = flops_after / flops_before
print("FLOPs ratio:", ratio)
net = net.to(device)
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
with torch.no_grad():
for index, sample in enumerate(tqdm(dataloader_train, leave=False)):
_ = net.get_loss(sample)
if index > 100:
break
strategy_score = net.get_eval_scores(dataloader_val)["accuracy"]
print("Result file:{}, Strategy ID:{}, Evaluation score:{}".format(
opt.search_result, opt.strategy_id, strategy_score))
##################### Fine-tuning #########################
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
net.optimizer, opt.epoch)
reporter = Reporter(opt)
best_acc = 0
net._net.train()
for epoch in range(1, opt.epoch + 1):
reporter.log_metric("lr", net.optimizer.param_groups[0]["lr"], epoch)
train_loss = train_epoch(
net,
dataloader_train,
net.optimizer,
)
reporter.log_metric("train_loss", train_loss, epoch)
lr_scheduler.step()
scores = net.get_eval_scores(dataloader_val)
print("==> Evaluation: Epoch={} Acc={}".format(epoch, str(scores)))
reporter.log_metric("eval_acc", scores["accuracy"], epoch)
if scores["accuracy"] > best_acc:
best_acc = scores["accuracy"]
reporter.log_metric("best_acc", best_acc, epoch)
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
print("==> Training epoch %d" % epoch)
def main(opt, channel_config, dataloader_train, dataloader_val, path):
# basic settings
torch.backends.cudnn.enabled = False
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = "cuda"
torch.backends.cudnn.benchmark = False
else:
device = "cpu"
##################### Get Dataloader ####################
# dummy_input is sample input of dataloaders
if hasattr(dataloader_val, "dataset"):
dummy_input = dataloader_val.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
load(net, path)
#net.load_checkpoint(opt.checkpoint)
##################### Load Pruning Strategy ###############
compression_scheduler = distiller.file_config(net.get_compress_part(),
net.optimizer,
opt.compress_schedule_path)
compression_scheduler = setCompressionScheduler(compression_scheduler,
channel_config)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=dummy_input)
flops_after, params_after = model_summary(net.get_compress_part(),
dummy_input)
net = net.to(device)
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
t = tqdm(dataloader_train, leave=False)
with torch.no_grad():
for index, sample in enumerate(t):
_ = net.get_loss(sample)
if index > 100:
break
strategy_score = net.get_eval_scores(dataloader_val)["accuracy"]
old = strategy_score
print("Evaluation score:{}".format(strategy_score))
##################### Fine-tuning #########################
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(net.optimizer,
100,
eta_min=5e-5)
#lr_scheduler=optim.lr_scheduler.StepLR(net.optimizer,5,0.9)
reporter = Reporter(opt)
best_acc = strategy_score
best_kappa = 0
net._net.train()
for epoch in range(1, opt.epoch + 1):
net.confusion_matrix.reset()
reporter.log_metric("lr", net.optimizer.param_groups[0]["lr"], epoch)
train_loss = train_epoch(
net,
dataloader_train,
net.optimizer,
)
reporter.log_metric("train_loss", train_loss, epoch)
lr_scheduler.step()
scores = net.get_eval_scores(dataloader_val)
kappa = CaluKappa(net.confusion_matrix)
print("==> Evaluation: Epoch={} Acc={}".format(epoch, str(scores)))
reporter.log_metric("eval_acc", scores["accuracy"], epoch)
reporter.log_metric("kappa", kappa, epoch)
if scores["accuracy"] > best_acc:
best_acc = scores["accuracy"]
best_kappa = kappa
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
reporter.log_metric("best_acc", best_acc, epoch)
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
print("==> Training epoch %d" % epoch)
"""将模型转换为torch script保存"""
ckpt_name = "{}_best.pth".format(opt.exp_name)
load(net, os.path.join(reporter.ckpt_log_dir, ckpt_name))
net._net.eval()
traced_script_module = torch.jit.trace(net._net,
torch.rand(1, 3, 256, 256))
traced_script_module.save(os.path.join(reporter.log_dir, "model.pt"))
del net
return old, best_acc, best_kappa, flops_after, params_after