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)