def decompose(self):
print('\n{}Bayesian Begin'.format('\033[33m'))
print('↓↓↓↓↓↓↓↓↓↓↓↓↓↓{}\n'.format('\033[0m'))
self.conv_decomposition()
print('-------------> Decomposition Finish')
print('Final Fine_tune ...')
fine_tune(self.model, 30)
fine_tune_test(self.model, self.testloader, True)
print('The Decomposed Model ...')
print(self.model)
mac, weight = self.get_model_mac_weight(self.model)
#deploy to target
save_model_name = export_onnx_model(self.model)
decomp_runtime = deploy_by_rpc(save_model_name)
self.decomp_runtime_ms = decomp_runtime * 1000
os.remove(save_model_name)
tmp_decomp_predict_runtime, tmp_decomp_layer_runtime = self.get_model_predict_runtime(self.model)
print('Origin_MAC: {}, Origin_Weight: {}, Origin_Runtime: {}, Origin_Predict_Runtime: {}'.format(self.origin_mac, self.origin_weight, self.real_model_runtime, self.origin_model_runtime))
print('Decomp_MAC: {}, Decomp_Weight: {}, Decomp_Runtime: {}, Decomp_Predict_Runtime: {}'.format(mac, weight, self.decomp_runtime_ms, tmp_decomp_predict_runtime))
print('Speedup : {}'.format(float(self.real_model_runtime/self.decomp_runtime_ms)))
def decompose(self):
print('\n{}Bayesian Begin'.format('\033[33m'))
print('↓↓↓↓↓↓↓↓↓↓↓↓↓↓{}\n'.format('\033[0m'))
bayesian_iter = self.conv_decomposition()
print('-------------> Decomposition Finish')
print('Final Fine_tune ...')
fine_tune(self.model, 30)
acc = fine_tune_test(self.model, self.testloader, True)
print('The Decomposed Model ...')
print(self.model)
mac, weight = self.get_model_mac_weight(self.model)
#deploy to target
save_model_name = export_onnx_model(self.model)
decomp_runtime = deploy_by_rpc(save_model_name)
self.decomp_runtime_ms = decomp_runtime * 1000
os.remove(save_model_name)
tmp_decomp_predict_runtime, tmp_decomp_layer_runtime = self.get_model_predict_runtime(self.model)
print('Origin_MAC: {}, Origin_Weight: {}, Origin_Runtime: {}, Origin_Predict_Runtime: {}'.format(self.origin_mac, self.origin_weight, self.real_model_runtime, self.origin_model_runtime))
print('Decomp_MAC: {}, Decomp_Weight: {}, Decomp_Runtime: {}, Decomp_Predict_Runtime: {}'.format(mac, weight, self.decomp_runtime_ms, tmp_decomp_predict_runtime))
print('Speedup : {}'.format(float(self.real_model_runtime/self.decomp_runtime_ms)))
state = {
'model': self.model,
'bayesian_iter': bayesian_iter,
'acc': acc,
'real_runtime': self.decomp_runtime_ms,
'predict_runtime': tmp_decomp_predict_runtime,
'mac': mac,
'weight': weight,
}
torch.save(state, 'result/decomposed/' + str(self.constrain) + '_alexnet_model')
def prune(self):
print('\n{}NetAdapt Begin{}'.format(bc.yellow, bc.end))
print('{}↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓{}'.format(bc.yellow, bc.end))
print("\n{}Runtime_budget is {}{}".format(bc.green, runtime_budget,
bc.end))
#Get the accuracy before prunning
self.test()
self.model.train()
#Make sure all the layers are trainable
for param in self.model.features.parameters():
param.requires_grad = True
model_image_size = {}
if (dataset == 'cifar'):
in_image_size = 32
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key],
torch.nn.modules.conv.Conv2d):
conv_layer = self.model.features._modules[key]
after_image_size = (
(in_image_size - conv_layer.kernel_size[0] + 2 *
conv_layer.padding[0]) // conv_layer.stride[0]) + 1
model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
elif isinstance(self.model.features._modules[key],
torch.nn.modules.MaxPool2d):
maxpool_layer = self.model.features._modules[key]
after_image_size = (
(in_image_size - maxpool_layer.kernel_size) //
maxpool_layer.stride) + 1
model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
print(model_image_size)
#Get the index of each layer
self.layer_index = []
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key],
torch.nn.modules.conv.Conv2d):
self.layer_index.append(key)
#Get the runtime before prunning
self.origin_predict_layer_runtime = {}
origin_predict_runtime = 0.0
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key],
torch.nn.modules.conv.Conv2d):
conv_layer = self.model.features._modules[key]
tmp_perf_runtime = self.perf_model.conv_predict(model_image_size[key][0], conv_layer.in_channels, conv_layer.out_channels ,\
conv_layer.kernel_size[0], conv_layer.stride[0], conv_layer.padding[0])
self.origin_predict_layer_runtime['conv_' +
key] = tmp_perf_runtime
origin_predict_runtime += tmp_perf_runtime
for i, key in enumerate(self.model.classifier._modules.keys()):
if isinstance(self.model.classifier._modules[key],
torch.nn.modules.linear.Linear):
fc_layer = self.model.classifier._modules[key]
tmp_perf_runtime = self.perf_model.fc_predict(
fc_layer.in_features, fc_layer.out_features)
self.origin_predict_layer_runtime['fc_' +
key] = tmp_perf_runtime
origin_predict_runtime += tmp_perf_runtime
print('{}Predict_Origin_predict_runtime{}: {}ms'.format(
bc.green, bc.end, origin_predict_runtime))
print('{}Predict_Origin_predict_layer_runtime{}: {}\n'.format(
bc.green, bc.end, self.origin_predict_layer_runtime))
self.decomposed_predict_layer_runtime = copy.deepcopy(
self.origin_predict_layer_runtime)
origin_model = copy.deepcopy(self.model)
#model_runtime = origin_runtime * 1000
perf_model_runtime = origin_predict_runtime
iteration_count = 1
while (perf_model_runtime > runtime_budget):
print('{}Iteration {}{}'.format(bc.red, iteration_count, bc.end))
print('{}--------------------------------------------{}'.format(
bc.red, bc.end))
if (iteration_count > 100):
import sys
print('iteration > 100')
sys.exit(1)
number_of_filters = self.total_num_filters()
print("Ranking filters.. ")
prune_targets = self.get_rank_to_prune(number_of_filters)
print('{}Number_of_filters{}: {}'.format(bc.green, bc.end,
number_of_filters))
print('{}Initialization{}: {}, {}Decay{}: {}'.format(
bc.green, bc.end, self.initialization, bc.green, bc.end,
self.decay_rate))
#print('Travis model: ', self.model)
#print('Travis prune_targets: ', prune_targets)
print('')
layer_record = {}
model_record = {}
#tmp_model = copy.deepcopy(self.model)
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key],
torch.nn.modules.conv.Conv2d):
prune_conv_layer = self.model.features._modules[key]
if (self.layer_index.index(key) !=
len(self.layer_index) - 1):
next_conv_layer = self.model.features._modules[
self.layer_index[self.layer_index.index(key) + 1]]
else:
next_conv_layer = None
print('{}Convolution Layer {}{}'.format(
bc.light_blue, key, bc.end))
print('Curr Layer: {}, Next Layer: {}'.format(
prune_conv_layer, next_conv_layer))
tmp_model = copy.deepcopy(self.model)
#print('Travis tmp_model: {}'.format(tmp_model))
layer_all_filter = []
for i in prune_targets:
if (str(i[0]) == key):
layer_all_filter.append(i)
tmp_perf_runtime = self.perf_model.conv_predict(model_image_size[key][0], prune_conv_layer.in_channels, prune_conv_layer.out_channels ,\
prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0], prune_conv_layer.padding[0])
if (tmp_perf_runtime <= self.initialization):
print('{}tmp_perf_runtime <= self.initialization{}'.
format(bc.purple, bc.end))
print('tmp_perf_runtime: {}, self.initialzation: {}'.
format(tmp_perf_runtime, self.initialization))
print('image_size: {}, in_channels: {}, out_channels: {}, kernel_size: {}, stride: {}, padding: {}'.format(model_image_size[key][0],\
prune_conv_layer.in_channels, prune_conv_layer.out_channels, prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0],\
prune_conv_layer.padding[0]))
layer_record[key] = [0, None, None]
print('')
continue
if (next_conv_layer != None):
prune_layer_runtime_1 = self.perf_model.conv_predict(model_image_size[key][0], prune_conv_layer.in_channels, prune_conv_layer.out_channels ,\
prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0], prune_conv_layer.padding[0])
prune_layer_runtime_2 = self.perf_model.conv_predict(model_image_size[self.layer_index[self.layer_index.index(key)+1]][0], next_conv_layer.in_channels, next_conv_layer.out_channels ,\
next_conv_layer.kernel_size[0], next_conv_layer.stride[0], next_conv_layer.padding[0])
prune_layer_runtime = prune_layer_runtime_1 + prune_layer_runtime_2
else:
prune_layer_runtime = self.perf_model.conv_predict(model_image_size[key][0], prune_conv_layer.in_channels, prune_conv_layer.out_channels ,\
prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0], prune_conv_layer.padding[0])
for tmp_out_channel in range(
prune_conv_layer.out_channels - 1, 1, -1):
if (next_conv_layer != None):
tmp_runtime_1 = self.perf_model.conv_predict(model_image_size[key][0], prune_conv_layer.in_channels, tmp_out_channel ,\
prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0], prune_conv_layer.padding[0])
tmp_runtime_2 = self.perf_model.conv_predict(model_image_size[self.layer_index[self.layer_index.index(key)+1]][0], tmp_out_channel, next_conv_layer.out_channels ,\
next_conv_layer.kernel_size[0], next_conv_layer.stride[0], next_conv_layer.padding[0])
tmp_runtime = tmp_runtime_1 + tmp_runtime_2
else:
tmp_runtime_1 = self.perf_model.conv_predict(model_image_size[key][0], prune_conv_layer.in_channels, tmp_out_channel ,\
prune_conv_layer.kernel_size[0], prune_conv_layer.stride[0], prune_conv_layer.padding[0])
tmp_runtime = tmp_runtime_1
#print('Travis tmp_out_channel: {}, tmp_runtime: {}, prune_layer_runtime: {}'.format(tmp_out_channel, tmp_runtime, prune_layer_runtime))
if ((prune_layer_runtime - tmp_runtime) >=
self.initialization):
print('Travis prune_layer_runtime: {}, tmp_runtime: {}, prune_layer_runtime-tmp_runtime: {}'.format(prune_layer_runtime, \
tmp_runtime, prune_layer_runtime-tmp_runtime))
num_filter_to_prune = prune_conv_layer.out_channels - tmp_out_channel
break
layer_prune_target = layer_all_filter[
0:num_filter_to_prune]
#for i in layer_prune_target:
# print(i)
prune_plan = self.prunner.get_layer_prunning_plan(
layer_prune_target)
#print(layer_prune_target)
#print(prune_plan)
layers_prunned = {}
for layer_index, filter_index in prune_plan:
if layer_index not in layers_prunned:
layers_prunned[layer_index] = 0
layers_prunned[
layer_index] = layers_prunned[layer_index] + 1
print("Layers that will be prunned", layers_prunned)
print("Prunning filters.. ")
cpu_model = tmp_model.cpu()
for layer_index, filter_index in prune_plan:
model = prune_vgg16_conv_layer(cpu_model, layer_index,
filter_index)
tmp_model = cpu_model.cuda()
tmp_model, tmp_acc = fine_tune(tmp_model, True)
layer_record[key] = [tmp_acc, prune_plan, tmp_model]
print('Acc after Fine_Tune {}'.format(tmp_acc))
print('')
#print('Travis Model', self.model)
acc_max = [0, -1]
for i, key in enumerate(layer_record.keys()):
if layer_record[key][0] > acc_max[1]:
acc_max = [key, layer_record[key][0]]
print('{}Pick max acc..{} key: {}, acc: {}'.format(
bc.blue, bc.end, acc_max[0], acc_max[1]))
## Travis Test
self.model = layer_record[acc_max[0]][2]
self.prunner.modify_model(layer_record[acc_max[0]][2])
print(self.model)
self.test()
tmp_latency = 0.0
tmp_layer_latency = {}
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key],
torch.nn.modules.conv.Conv2d):
tmp_conv_layer = self.model.features._modules[key]
tmp_perf_runtime = self.perf_model.conv_predict(model_image_size[key][0], tmp_conv_layer.in_channels, tmp_conv_layer.out_channels ,\
tmp_conv_layer.kernel_size[0], tmp_conv_layer.stride[0], tmp_conv_layer.padding[0])
tmp_layer_latency['conv_' + key] = tmp_perf_runtime
tmp_latency += tmp_perf_runtime
for i, key in enumerate(self.model.classifier._modules.keys()):
if isinstance(self.model.classifier._modules[key],
torch.nn.modules.linear.Linear):
fc_layer = self.model.classifier._modules[key]
tmp_perf_runtime = self.perf_model.fc_predict(
fc_layer.in_features, fc_layer.out_features)
tmp_layer_latency['fc_' + key] = tmp_perf_runtime
tmp_latency += tmp_perf_runtime
print(
'{}Predict Runtime after iteration {}: {}ms, reduction: {}ms{}'
.format('\033[32m', iteration_count, tmp_latency,
perf_model_runtime - tmp_latency, '\033[0m'))
print('Runtime for each layer: {}'.format(tmp_layer_latency))
perf_model_runtime = tmp_latency
#torch.save(self.model, "result/prunned_model/model_" + str(runtime_budget) + '_' + str(iteration_count) + "_prunned")
## Get runtime after one iteration
tmp_save_model_name = export_onnx_model(self.model)
tmp_model_runtime = deploy_by_rpc(tmp_save_model_name) * 1000
print('{}Real Runtime after iteration {}: {}ms{}'.format(
'\033[32m', iteration_count, tmp_model_runtime, '\033[0m'))
os.remove(tmp_save_model_name)
iteration_count += 1
self.initialization *= self.decay_rate
print('')
print("{}Finished. Going to fine tune the model a bit more{}".format(
'\033[33m', '\033[0m'))
optimizer = optim.SGD(self.model.parameters(), lr=0.0001, momentum=0.9)
self.train(optimizer, epoches=15)
print('Travis self.model', self.model)
final_acc = fine_tune_test(self.model, self.test_data_loader, True)
## Get runtime after one iteration
tmp_save_model_name = export_onnx_model(self.model)
tmp_model_runtime = deploy_by_rpc(tmp_save_model_name) * 1000
print('Runtime after pruning: {}ms'.format(tmp_model_runtime))
os.remove(tmp_save_model_name)
model_runtime_after_pruning = tmp_model_runtime
state = {
'net': self.model,
'acc': final_acc,
'iteration_count': iteration_count,
'model_runtime': model_runtime_after_pruning,
}
torch.save(
state,
"result/prunned_model/model_" + str(runtime_budget) + "_prunned")
def __init__(self, **kwargs):
opt.parse(kwargs)
tl.set_backend('pytorch')
self.dataset = "cifar"
self.decomposed_layer_info = {'key': -1, 'image_size': -1, 'kernel_size': -1, 'stride': -1, 'padding': -1}
self.layer_budget = {}
self.origin_layer_runtime = {}
self.origin_model_runtime = 0.0
self.VBMF_layer_rank = {}
self.constrain = opt.constrain
self.conv_target_rate = 0.0
self.fc_target_rate = 0.0
self.user_budget = 1
self.real_model_runtime = 0.0
self.remain_budget = 0.0
self.origin_model_constrain = 0.0
self.search_runtime = {}
self.bayesian_iter = {}
# Configure Logger
self.logger = logging.getLogger()
log_file = logging.FileHandler('result/log/test.log')
self.logger.addHandler(log_file)
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
log_file.setFormatter(formatter)
self.logger.setLevel(logging.DEBUG)
# Load Perf Model
self.perf_model = Estimator()
# Load Pre-trained Model
if(opt.load_model_path is None):
import sys
print('set the model path')
sys.exit(-1)
else:
checkpoint = torch.load(opt.load_model_path)
if(type(checkpoint) is dict):
checkpoint = checkpoint['net']
self.model = checkpoint.cuda()
print(self.model)
# Preparing data
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)),
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform_train)
self.trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform_test)
self.testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4)
print('\n{}Model Info'.format('\033[33m'))
print('↓↓↓↓↓↓↓↓↓↓↓↓↓↓{}\n'.format('\033[0m'))
# Set Criterion
self.criterion = torch.nn.CrossEntropyLoss()
# Calculate Image_size for each layer
self.model_image_size = {}
if(self.dataset == 'cifar'):
in_image_size = 32
for i, key in enumerate(self.model.features._modules.keys()):
if isinstance(self.model.features._modules[key], torch.nn.modules.conv.Conv2d):
conv_layer = self.model.features._modules[key]
after_image_size = ((in_image_size - conv_layer.kernel_size[0] + 2*conv_layer.padding[0]) // conv_layer.stride[0] )+ 1
self.model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
elif isinstance(self.model.features._modules[key], torch.nn.modules.MaxPool2d):
maxpool_layer = self.model.features._modules[key]
after_image_size = ((in_image_size - maxpool_layer.kernel_size) // maxpool_layer.stride )+ 1
self.model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
print('{}Image_Size{}: {}'.format('\033[36m', '\033[0m', self.model_image_size))
# Get Origin MAC and Weight and runtime
self.origin_mac, self.origin_weight = self.get_model_mac_weight(self.model)
self.origin_model_runtime, self.origin_layer_runtime = self.get_model_predict_runtime(self.model)
self.origin_model_constrain, _ = self.get_model_predict_runtime_without_small(self.model)
#print('self.origin_model_runtime: {}, self.get_model_predict_runtime: {}'.format(self.origin_model_runtime, self.get_model_predict_runtime(self.model)))
#deploy to target
save_model_name = export_onnx_model(self.model)
decomp_runtime = deploy_by_rpc(save_model_name)
self.real_model_runtime = decomp_runtime * 1000
os.remove(save_model_name)
print('{}Origin_MAC{}: {}, {}Origin_Weight{}: {}'.format('\033[36m', '\033[0m', self.origin_mac, '\033[36m', '\033[0m', self.origin_weight))
#print('{}Origin_Weight{}: {}'.format('\033[36m', '\033[0m', self.origin_weight))
print('{}Pred_Origin_Runtime{}: {}, {}Real_Origin_Runtime{}: {}'.format('\033[36m', '\033[0m', self.origin_model_runtime, '\033[36m', '\033[0m', self.real_model_runtime))
#print('{}Real_Origin_Runtime{}: {}'.format('\033[36m', '\033[0m', self.real_model_runtime))
print('{}Origin_Layer_Runtime{}: {}'.format('\033[36m', '\033[0m', self.origin_layer_runtime))
print('{}Origin_Model_Constrain{}: {}'.format('\033[36m', '\033[0m', self.origin_model_constrain))
self.VBMF_layer_rank = self.get_VBMF_layer_rank()
if(self.constrain > 0):
# Calculate importance for each layer
self.layer_importance = self.get_layer_importance()
print('{}Layer Importance{}: {}'.format('\033[36m', '\033[0m', self.layer_importance))
# Get Layer Budget
self.layer_budget = self.get_layer_budget()
def conv_decomposition(self):
VBMF_model = torch.load('checkpoint/VBMF_alexnet_model')
_, tmp_VBMF_layer_runtime = self.get_model_predict_runtime(VBMF_model)
bayesian_iter = {}
#remain_budget = 0.0
for i, key in enumerate(self.model.features._modules.keys()):
if(i == 0):
continue
if isinstance(self.model.features._modules[key], torch.nn.modules.conv.Conv2d):
conv_layer_to_decompose = self.model.features._modules[key]
#print('\nTravis conv_layer_to_decompose rank1:{}, rank2:{}'.format(conv_layer_to_decompose.in_channels, conv_layer_to_decompose.out_channels))
print('\n{}Layer Info{}'.format('\033[31m', '\033[0m'))
print('{}-----------------------------------------{}'.format('\033[94m', '\033[0m'))
print('Rank1:{}, Rank2:{}'.format(conv_layer_to_decompose.in_channels, conv_layer_to_decompose.out_channels))
#print(self.estimate(conv_layer_to_decompose, key))
self.decomposed_layer_info['key'] = key
self.decomposed_layer_info['image_size'] = self.model_image_size[key][0]
self.decomposed_layer_info['kernel_size'] = conv_layer_to_decompose.kernel_size[0]
self.decomposed_layer_info['stride'] = conv_layer_to_decompose.stride[0]
self.decomposed_layer_info['padding'] = conv_layer_to_decompose.padding[0]
self.conv_target_rate = 0.0
iteration_count = 0
while(True):
print('{}Iteration{}'.format('\033[31m', '\033[0m'))
print('{}-----------------------------------------{}'.format('\033[94m', '\033[0m'))
print('Iteration {}{}{} Target_Rate: {}{}{}'.format('\033[32m', iteration_count, '\033[0m', '\033[32m', self.conv_target_rate, '\033[0m'))
iteration_count += 1
print('Travis key: {}, search_runtime: {}, VBMF_layer_runtime: {}'.format(key, self.search_runtime['conv_'+key], tmp_VBMF_layer_runtime))
if(self.search_runtime['conv_'+key] == tmp_VBMF_layer_runtime['conv_'+key]):
ranks = [self.VBMF_layer_rank['conv_'+key][0], self.VBMF_layer_rank['conv_'+key][1]]
b_iter = 0
else:
ranks, b_iter = self.conv_bayesian([conv_layer_to_decompose.in_channels, conv_layer_to_decompose.out_channels])
bayesian_iter['conv_'+key] = b_iter
tmp_ms_1 = self.estimate_with_config([self.decomposed_layer_info['image_size'], \
conv_layer_to_decompose.in_channels, ranks[0], \
1, \
self.decomposed_layer_info['stride'], \
self.decomposed_layer_info['padding']])
tmp_ms_2 = self.estimate_with_config([self.decomposed_layer_info['image_size'], \
ranks[0], ranks[1], \
self.decomposed_layer_info['kernel_size'], \
self.decomposed_layer_info['stride'], \
self.decomposed_layer_info['padding']])
tmp_ms_3 = self.estimate_with_config([self.decomposed_layer_info['image_size'], \
ranks[1], conv_layer_to_decompose.out_channels, \
1, \
self.decomposed_layer_info['stride'], \
self.decomposed_layer_info['padding']])
tmp_runtime_ms = tmp_ms_1 + tmp_ms_2 + tmp_ms_3
print('Travis tmp_ms_1: {}, tmp_ms_2: {}, tmp_ms_3: {}'.format(tmp_ms_1, tmp_ms_2, tmp_ms_3))
if(self.constrain == 0):
break
else:
if((tmp_runtime_ms) <= (self.search_runtime['conv_'+ key] + self.remain_budget)):
print('Travis Constrain: {}, Search_runtime: {}, Layer_budget: {}, Remain_budget: {}'.format(self.constrain, \
tmp_runtime_ms, \
self.search_runtime['conv_'+key], self.remain_budget))
print('Travis Layer_budget + Remain_budget: {}'.format(self.search_runtime['conv_'+key]+self.remain_budget))
self.remain_budget = (self.search_runtime['conv_'+key] + self.remain_budget) - (tmp_runtime_ms)
assert self.remain_budget >= 0
print('Updated Remain Budget: {}'.format(self.remain_budget))
# prevent TVM from disconnecting
save_model_name = export_onnx_model(self.model)
decomp_runtime = deploy_by_rpc(save_model_name)
self.decomp_runtime_ms = decomp_runtime * 1000
os.remove(save_model_name)
break
else:
print('Update the objective function ...')
self.conv_target_rate += 0.1
print('Travis Constrain: {}, Search_runtime: {}, Layer_budget: {}, Remain_budget: {}'.format(self.constrain, \
tmp_runtime_ms, \
self.search_runtime['conv_'+key], self.remain_budget))
print('Travis Layer_budget + Remain_budget: {}'.format(self.search_runtime['conv_'+key]+self.remain_budget))
print('{}Fine Tune{}'.format('\033[31m', '\033[0m'))
print('{}-----------------------------------------{}'.format('\033[94m', '\033[0m'))
ranks[0], ranks[1] = ranks[1], ranks[0]
decompose = tucker_decomposition_conv_layer_without_rank(self.model.features._modules[key],ranks)
self.model.features._modules[key] = decompose
fine_tune(self.model, 10)
fine_tune_test(self.model, self.testloader, True)
return bayesian_iter
def print_model(**kwargs):
opt.parse(kwargs)
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform_test)
testloader = torch.utils.data.DataLoader(testset,
batch_size=100,
shuffle=False,
num_workers=4)
if opt.load_model_path:
model = torch.load(opt.load_model_path)
print(model)
else:
import sys
print('set the load_model_path')
sys.exit(1)
if (type(model) is dict):
model = model['net']
model = model.cuda()
#model, acc = fine_tune(model, True)
best_acc = fine_tune_test(model, testloader, True)
print(best_acc)
#Get the runtime before prunning
tmp_save_model_name = export_onnx_model(model)
origin_runtime = deploy_by_rpc(tmp_save_model_name)
print('Real Runtime: {}ms'.format(origin_runtime * 1000))
os.remove(tmp_save_model_name)
#total_mac, total_weight = get_model_mac_weight(model)
#print('Mac: {}, Weight: {}'.format(total_mac, total_weight))
# Calculate Image_size for each layer
#model_image_size = {'0': [32, 16], '2': [16, 8], '3': [8, 8], '5': [8, 4], '6': [4, 4], '8': [4, 4], '10': [4, 4], '12': [4, 2]}
model_image_size = {}
in_image_size = 32
for i, key in enumerate(model.features._modules.keys()):
if isinstance(model.features._modules[key],
torch.nn.modules.conv.Conv2d):
conv_layer = model.features._modules[key]
after_image_size = (
(in_image_size - conv_layer.kernel_size[0] +
2 * conv_layer.padding[0]) // conv_layer.stride[0]) + 1
model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
elif isinstance(model.features._modules[key],
torch.nn.modules.MaxPool2d):
maxpool_layer = model.features._modules[key]
after_image_size = ((in_image_size - maxpool_layer.kernel_size) //
maxpool_layer.stride) + 1
model_image_size[key] = [in_image_size, after_image_size]
in_image_size = after_image_size
print('{}Image_Size{}: {}'.format('\033[36m', '\033[0m', model_image_size))
#Get the Predicted Runtime
predicted_model_runtime, _ = get_model_runtime(model, model_image_size)
print('Predicted Runtime: {}ms'.format(predicted_model_runtime))