def trans_model(model, export_model_dir, model_version): """ h5 model to pb model :param model: :param export_model_dir: :param model_version: :return: """ with tf.get_default_graph().as_default(): # prediction_signature tensor_info_input_0 = tf.saved_model.utils.build_tensor_info( model.input[0]) tensor_info_input_1 = tf.saved_model.utils.build_tensor_info( model.input[1]) tensor_info_input_2 = tf.saved_model.utils.build_tensor_info( model.input[2]) tensor_info_input_3 = tf.saved_model.utils.build_tensor_info( model.input[3]) tensor_info_output = tf.saved_model.utils.build_tensor_info( model.output) print(model.input) print(model.output.shape, '**', tensor_info_output) prediction_signature = ( tf.saved_model.signature_def_utils.build_signature_def( inputs={ 'input_0': tensor_info_input_0, 'input_1': tensor_info_input_1, 'input_2': tensor_info_input_2, 'input_3': tensor_info_input_3 }, # Tensorflow.TensorInfo outputs={'result': tensor_info_output}, # method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME) method_name="tensorflow/serving/predict")) print('-----prediction_signature created successfully-----') os.mkdir(export_model_dir) export_path_base = export_model_dir export_path = os.path.join(tf.compat.as_bytes(export_path_base), tf.compat.as_bytes(str(model_version))) builder = tf.saved_model.builder.SavedModelBuilder(export_path) builder.add_meta_graph_and_variables( sess=K.get_session(), tags=[tf.saved_model.tag_constants.SERVING], signature_def_map={ 'predict': prediction_signature, tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature, }, ) print('Export path(%s) ready to export trained model' % export_path, '\n starting to export model...') # builder.save(as_text=True) builder.save() print('Done exporting!')
def get_num_gpus(with_keras=False): if with_keras: from keras import get_session return len(get_session().list_devices()) host_gpus = dict(gibson='0,1', valis='', mira='1,2') host_gpu = host_gpus.get(hostname(), '0') os.environ.setdefault('CUDA_VISIBLE_DEVICES', host_gpu) cuda_env = os.environ.get('CUDA_VISIBLE_DEVICES', host_gpu) num_gpus = 0 if cuda_env == '' else len(cuda_env.split(',')) return num_gpus
def auc(y_true, y_pred): auc = tf.metrics.auc(y_true, y_pred)[1] K.get_session().run(tf.local_variables_initializer()) return auc