def __init__(self, remote_vars, local_vars, *args, **kwargs): if local_vars: self.has_init = tf.is_variable_initialized(local_vars[0]) else: self.has_init = tf.constant(True) ready_for_local_init_op = tf.report_uninitialized_variables( var_list=remote_vars) self.init_op = tf.group(tf.initialize_variables(local_vars), *tf.get_collection(tf.GraphKeys.LOCAL_INIT_OP)) if "scaffold" in kwargs: # TODO(lmetz) I think this could technically be supported? raise ValueError("Do not set scaffold on the session creator.") scaffold = tf.train.Scaffold( ready_for_local_init_op=ready_for_local_init_op, ready_op=ready_for_local_init_op, local_init_op=ready_for_local_init_op, init_op=ready_for_local_init_op, init_fn=self._maybe_initialize_local_vars_and_state_fn, summary_op=tf.summary.merge([tf.summary.scalar("dummy", 0)])) kwargs["scaffold"] = scaffold super(WorkerSessionCreator, self).__init__(*args, **kwargs)
def scaffold_fn(): return tf.train.Scaffold( local_init_op=tf.group( tf.train.Scaffold.default_local_init_op(), lowering.copy_masters_to_slices(), name="mtf_local_init_op"), ready_op=tf.concat( [tf.report_uninitialized_variables(), resources.report_uninitialized_resources()], axis=0, name="mtf_ready_op"))
def test_restore_map_for_classification_ckpt(self): # TODO(rathodv): Support TF2.X if self.is_tf2(): return # Define mock tensorflow classification graph and save variables. test_graph_classification = tf.Graph() with test_graph_classification.as_default(): image = tf.placeholder(dtype=tf.float32, shape=[1, 20, 20, 3]) with tf.variable_scope('mock_model'): net = slim.conv2d(image, num_outputs=32, kernel_size=1, scope='layer1') slim.conv2d(net, num_outputs=3, kernel_size=1, scope='layer2') init_op = tf.global_variables_initializer() saver = tf.train.Saver() save_path = self.get_temp_dir() with self.session(graph=test_graph_classification) as sess: sess.run(init_op) saved_model_path = saver.save(sess, save_path) # Create tensorflow detection graph and load variables from # classification checkpoint. test_graph_detection = tf.Graph() with test_graph_detection.as_default(): model, _, _, _ = self._create_model() inputs_shape = [2, 2, 2, 3] inputs = tf.cast(tf.random_uniform(inputs_shape, minval=0, maxval=255, dtype=tf.int32), dtype=tf.float32) preprocessed_inputs, true_image_shapes = model.preprocess(inputs) prediction_dict = model.predict(preprocessed_inputs, true_image_shapes) model.postprocess(prediction_dict, true_image_shapes) another_variable = tf.Variable([17.0], name='another_variable') # pylint: disable=unused-variable var_map = model.restore_map( fine_tune_checkpoint_type='classification') self.assertNotIn('another_variable', var_map) self.assertIsInstance(var_map, dict) saver = tf.train.Saver(var_map) with self.session(graph=test_graph_detection) as sess: saver.restore(sess, saved_model_path) for var in sess.run(tf.report_uninitialized_variables()): self.assertNotIn(six.ensure_binary('FeatureExtractor'), var)
def freeze_checkpoints(graph_def: tf_v1.GraphDef, checkpoint_dir: str, output_node_names: list): """ Loads all the variables in a graph and stores them in a separate dictionary. Freezes output nodes in the graph :param graph_def: GraphDef object holding the network. :param checkpoint_dir: path to directory with checkpoint files with values of graph variables. :param output_node_names: list of output node names. :return: GraphDef containing a simplified version of the original. """ log.debug("Loading checkpoint files from directory: {}".format(checkpoint_dir)) checkpoint_files = [] for checkpoint_name in sorted(os.listdir(checkpoint_dir)): checkpoint_path = os.path.join(checkpoint_dir, checkpoint_name) if os.path.isfile(checkpoint_path): checkpoint_files.append(checkpoint_path) log.debug("File {} will be loaded".format(checkpoint_path)) else: log.debug("Path {} is not a file. Skipping") if len(checkpoint_files) == 0: raise Error("There are no checkpoint files in directory: {}".format(checkpoint_dir)) tf_v1.import_graph_def(graph_def, name='') with tf_v1.Session() as sess: uninitialized_variables = [str(v, 'utf-8') for v in set(sess.run(tf_v1.report_uninitialized_variables()))] all_variables = [n.name for n in sess.graph.as_graph_def().node if n.op in ['Variable', 'VariableV2']] white_list = [v for v in all_variables if v not in uninitialized_variables] black_list = [v for v in all_variables if v in uninitialized_variables] output_graph_def = tf_v1.graph_util.convert_variables_to_constants(sess, graph_def, output_node_names, variable_names_whitelist=white_list, variable_names_blacklist=black_list) variable_values = {} for checkpoint_file in checkpoint_files: log.debug("Loading {}".format(checkpoint_file)) with tf_v1.Session() as sess: var_list = {} var_to_shape_map = tf_v1.train.load_checkpoint(checkpoint_file).get_variable_to_shape_map() for key in var_to_shape_map: try: tensor = sess.graph.get_operation_by_name(key).outputs[0] except KeyError: continue var_list[key] = tensor tf_v1.train.Saver(var_list=var_list).restore(sess, checkpoint_file) for name, tensor in var_list.items(): variable_values[name] = sess.run(tensor) return output_graph_def, variable_values
def test_restore_map_for_detection_ckpt(self): # TODO(rathodv): Support TF2.X if self.is_tf2(): return model, _, _, _ = self._create_model() model.predict(tf.constant( np.array([[[[0, 0], [1, 1]], [[1, 0], [0, 1]]]], dtype=np.float32)), true_image_shapes=None) init_op = tf.global_variables_initializer() saver = tf.train.Saver() save_path = self.get_temp_dir() with self.session() as sess: sess.run(init_op) saved_model_path = saver.save(sess, save_path) var_map = model.restore_map( fine_tune_checkpoint_type='detection', load_all_detection_checkpoint_vars=False) self.assertIsInstance(var_map, dict) saver = tf.train.Saver(var_map) saver.restore(sess, saved_model_path) for var in sess.run(tf.report_uninitialized_variables()): self.assertNotIn('FeatureExtractor', var)
def tf_assert_all_init(sess): uninit_vars = sess.run(tf.report_uninitialized_variables()) assert len( uninit_vars ) == 0, 'Expected all variables to have been initialized, but these have not been: %s' % uninit_vars
def __init__(self, graph_path, target_size=(320, 240), tf_config=None, trt_bool=False, verbose=False): self.target_size = target_size # load graph logger.info('loading graph from %s(default size=%dx%d)' % (graph_path, target_size[0], target_size[1])) with tf.io.gfile.GFile(graph_path, 'rb') as f: graph_def = tf.compat.v1.GraphDef() graph_def.ParseFromString(f.read()) if trt_bool is True: output_nodes = ["Openpose/concat_stage7"] graph_def = trt.create_inference_graph( graph_def, output_nodes, max_batch_size=1, max_workspace_size_bytes=1 << 20, precision_mode="FP16", # precision_mode="INT8", minimum_segment_size=3, is_dynamic_op=True, maximum_cached_engines=int(1e3), use_calibration=True, ) self.graph = tf.compat.v1.get_default_graph() tf.import_graph_def(graph_def, name='TfPoseEstimator') self.persistent_sess = tf.compat.v1.Session(graph=self.graph, config=tf_config) if verbose: for ts in [ n.name for n in tf.compat.v1.get_default_graph().as_graph_def().node ]: print(ts) self.tensor_image = self.graph.get_tensor_by_name( 'TfPoseEstimator/image:0') self.tensor_output = self.graph.get_tensor_by_name( 'TfPoseEstimator/Openpose/concat_stage7:0') self.tensor_heatMat = self.tensor_output[:, :, :, :19] self.tensor_pafMat = self.tensor_output[:, :, :, 19:] self.upsample_size = tf.placeholder(dtype=tf.int32, shape=(2, ), name='upsample_size') self.tensor_heatMat_up = tf.image.resize_area( self.tensor_output[:, :, :, :19], self.upsample_size, align_corners=False, name='upsample_heatmat') self.tensor_pafMat_up = tf.image.resize_area( self.tensor_output[:, :, :, 19:], self.upsample_size, align_corners=False, name='upsample_pafmat') if trt_bool is True: smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0, 19) else: smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0) gaussian_heatMat = smoother.get_output() max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat, window_shape=(3, 3), pooling_type='MAX', padding='SAME') self.tensor_peaks = tf.where( tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat, tf.zeros_like(gaussian_heatMat)) self.heatMat = self.pafMat = None # warm-up self.persistent_sess.run( tf.compat.v1.variables_initializer([ v for v in tf.global_variables() if v.name.split(':')[0] in [ x.decode('utf-8') for x in self.persistent_sess.run( tf.report_uninitialized_variables()) ] ])) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [ np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32) ], self.upsample_size: [target_size[1], target_size[0]] }) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [ np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32) ], self.upsample_size: [target_size[1] // 2, target_size[0] // 2] }) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [ np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32) ], self.upsample_size: [target_size[1] // 4, target_size[0] // 4] }) # logs if self.tensor_image.dtype == tf.quint8: logger.info('quantization mode enabled.')
def __init__(self, graph_path, target_size=(320, 240), tf_config=None): self.target_size = target_size # load graph logger.info('loading graph from %s(default size=%dx%d)' % (graph_path, target_size[0], target_size[1])) with tf.gfile.GFile(graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) self.graph = tf.get_default_graph() tf.import_graph_def(graph_def, name='TfPoseEstimator') self.persistent_sess = tf.Session(graph=self.graph, config=tf_config) # for op in self.graph.get_operations(): # print(op.name) # for ts in [n.name for n in tf.get_default_graph().as_graph_def().node]: # print(ts) self.tensor_image = self.graph.get_tensor_by_name('TfPoseEstimator/image:0') self.tensor_output = self.graph.get_tensor_by_name('TfPoseEstimator/Openpose/concat_stage7:0') self.tensor_heatMat = self.tensor_output[:, :, :, :19] self.tensor_pafMat = self.tensor_output[:, :, :, 19:] self.upsample_size = tf.placeholder(dtype=tf.int32, shape=(2,), name='upsample_size') self.tensor_heatMat_up = tf.image.resize_area(self.tensor_output[:, :, :, :19], self.upsample_size, align_corners=False, name='upsample_heatmat') self.tensor_pafMat_up = tf.image.resize_area(self.tensor_output[:, :, :, 19:], self.upsample_size, align_corners=False, name='upsample_pafmat') smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0) gaussian_heatMat = smoother.get_output() max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat, window_shape=(3, 3), pooling_type='MAX', padding='SAME') self.tensor_peaks = tf.where(tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat, tf.zeros_like(gaussian_heatMat)) self.heatMat = self.pafMat = None # warm-up self.persistent_sess.run(tf.variables_initializer( [v for v in tf.global_variables() if v.name.split(':')[0] in [x.decode('utf-8') for x in self.persistent_sess.run(tf.report_uninitialized_variables())] ]) ) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32)], self.upsample_size: [target_size[1], target_size[0]] } ) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32)], self.upsample_size: [target_size[1] // 2, target_size[0] // 2] } ) self.persistent_sess.run( [self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up], feed_dict={ self.tensor_image: [np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32)], self.upsample_size: [target_size[1] // 4, target_size[0] // 4] } )