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