def split_original_model(model, input_shape, device, config, sensor_device,
edge_device, partition_idx, head_output_file_path,
tail_output_file_path, require_test, spbit):
print('Splitting an original DNN model')
modules = list()
z = torch.rand(1, *input_shape).to(device)
module_util.extract_decomposable_modules(model, z, modules)
head_module_list = list()
tail_module_list = list()
if partition_idx < 0:
teacher_model_config = config['teacher_model']
start_idx = teacher_model_config['start_idx']
end_idx = teacher_model_config['end_idx']
head_module_list.extend(modules[start_idx:end_idx])
tail_module_list.extend(modules[end_idx:])
else:
head_module_list.extend(modules[:partition_idx])
tail_module_list.extend(modules[partition_idx:])
head_network = nn.Sequential(*head_module_list)
tail_network = mimic_util.get_tail_network(config, tail_module_list)
file_util.save_pickle(head_network.to(sensor_device),
head_output_file_path)
file_util.save_pickle(tail_network.to(edge_device), tail_output_file_path)
if require_test:
test_split_model(model, head_network, tail_network, sensor_device,
edge_device, spbit, config)
def compute_ae_bottleneck_size(self, x, print_info=False):
z = self.head_model(x)
modules = list()
module_util.extract_decomposable_modules(self.autoencoder, z, modules)
modules = [module.to(x.device) for module in modules]
org_size = np.prod(x.size())
min_rate = None
bo = None
bqo = None
for i in range(len(modules)):
if isinstance(modules[i], nn.Linear):
z = z.view(z.size(0), -1)
z = modules[i](z)
rate = np.prod(z.size()) / org_size
if min_rate is None or rate < min_rate:
min_rate = rate
bo = pickle.dumps(z)
bqo = pickle.dumps(tensor_util.quantize_tensor(z))
output_data_size = sys.getsizeof(bo) / 1024
quantized_output_data_size = sys.getsizeof(bqo) / 1024
if print_info:
print(
'[Autoencoder bottleneck]\tScaled output size: {} [%]\tOutput data size: {} [KB]'
'\tQuantized output data size: {} [KB]'.format(
min_rate * 100.0, output_data_size,
quantized_output_data_size))
# Scaled bottleneck size, bottleneck data size [KB], Quantized bottleneck data size [KB]
return min_rate, output_data_size, quantized_output_data_size
def get_mimic_model(config,
org_model,
teacher_model_type,
teacher_model_config,
device,
head_model=None):
target_model = org_model.module if isinstance(
org_model, nn.DataParallel) else org_model
student_model =\
load_student_model(config, teacher_model_type, device) if head_model is None else head_model.to(device)
org_modules = list()
input_batch = torch.rand(config['input_shape']).unsqueeze(0).to(device)
module_util.extract_decomposable_modules(target_model, input_batch,
org_modules)
end_idx = teacher_model_config['end_idx']
tail_modules = org_modules[end_idx:]
mimic_model_config = config['mimic_model']
mimic_type = mimic_model_config['type']
if mimic_type.startswith('densenet'):
mimic_model = DenseNetMimic(student_model, tail_modules)
elif mimic_type.startswith('inception'):
mimic_model = InceptionMimic(student_model, tail_modules)
elif mimic_type.startswith('resnet'):
mimic_model = ResNetMimic(student_model, tail_modules)
elif mimic_type.startswith('mobilenet'):
mimic_model = MobileNetMimic(student_model, tail_modules)
else:
raise ValueError('mimic_type `{}` is not expected'.format(mimic_type))
return mimic_model.to(device)
def extract_teacher_model(model, input_shape, device, teacher_model_config):
modules = list()
module = model.module if isinstance(model, nn.DataParallel) else model
module_util.extract_decomposable_modules(
module,
torch.rand(1, *input_shape).to(device), modules)
start_idx = teacher_model_config['start_idx']
end_idx = teacher_model_config['end_idx']
return nn.Sequential(*modules[start_idx:end_idx]).to(device)
def extract_head_model(model, input_shape, device, partition_idx):
if partition_idx is None or partition_idx == 0:
return nn.Sequential()
modules = list()
module = model.module if isinstance(model,
(DataParallel,
DistributedDataParallel)) else model
module_util.extract_decomposable_modules(
module,
torch.rand(1, *input_shape).to(device), modules)
return nn.Sequential(*modules[:partition_idx]).to(device)
def extend_model(autoencoder, model, input_shape, device, partition_idx,
skip_bottleneck_size):
if partition_idx is None or partition_idx == 0:
return nn.Sequential(autoencoder, model)
modules = list()
module = model.module if isinstance(model,
(DataParallel,
DistributedDataParallel)) else model
x = torch.rand(1, *input_shape).to(device)
module_util.extract_decomposable_modules(module, x, modules)
extended_model = BaseExtendedModel(modules[:partition_idx], autoencoder,
modules[partition_idx:]).to(device)
if not skip_bottleneck_size:
extended_model.compute_ae_bottleneck_size(x, True)
return extended_model
def split_within_student_model(model, input_shape, device, config,
teacher_model_type, sensor_device, edge_device,
partition_idx, head_output_file_path,
tail_output_file_path, require_test, spbit):
print('Splitting within a student DNN model')
org_modules = list()
z = torch.rand(1, *input_shape).to(device)
plain_model = model.module if isinstance(
model, (DataParallel, DistributedDataParallel)) else model
module_util.extract_decomposable_modules(plain_model, z, org_modules)
student_model = mimic_util.load_student_model(config, teacher_model_type,
device)
student_modules = list()
module_util.extract_decomposable_modules(student_model, z, student_modules)
head_module_list = list()
tail_module_list = list()
teacher_model_config = config['teacher_model']
end_idx = teacher_model_config['end_idx']
if partition_idx < 0:
head_module_list.extend(student_modules)
else:
head_module_list.extend(student_modules[:partition_idx])
tail_module_list.extend(student_modules[partition_idx:])
tail_module_list.extend(org_modules[end_idx:])
head_network = nn.Sequential(*head_module_list)
tail_network = mimic_util.get_tail_network(config, tail_module_list)
file_util.save_pickle(head_network.to(sensor_device),
head_output_file_path)
file_util.save_pickle(tail_network.to(edge_device), tail_output_file_path)
if require_test:
device = torch.device(
'cuda' if next(model.parameters()).is_cuda else 'cpu')
mimic_model = mimic_util.get_mimic_model(config, model,
teacher_model_type,
teacher_model_config, device)
test_split_model(mimic_model, head_network, tail_network,
sensor_device, edge_device, spbit, config)