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 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))