def from_session(sess, inputs, outputs,
generate_backward=False,
allow_non_differentiable_input=True,
tf_session_config=None):
"""
Create a TFNet from an a session and the inputs and outpus endpoints
of the TensorFlow graph.
:param sess: the TensorFlow session contain all the variables
:param inputs: a list of TensorFlow Tensor represents the input endpoints
of the TensorFlow graph
:param outputs: a list of TensorFlow Tensor represents the output endpoints
of the TensorFlow graph
:param generate_backward: whether to generated a the backward graph, set true
if you want to train this TFNet
:param allow_non_differentiable_input: if set to yes, when input are not differentiable,
the gradient will be set to zero. if set to false, an error will be thrown.
:param tf_session_config: an optional tf.ConfigProto object to
set the session config in java side.
This config does not necessarily be the same with your current session.
E.g. sess_config = tf.ConfigProto(inter_op_parallelism_threads=1,
intra_op_parallelism_threads=1)
net = TFNet.from_session(sess, inputs, outputs, sess_config)
:return a TFNet
"""
from zoo.util.tf import export_tf
temp = tempfile.mkdtemp()
try:
export_tf(sess, temp, inputs, outputs,
generate_backward, allow_non_differentiable_input)
net = TFNet.from_export_folder(temp, tf_session_config)
finally:
import shutil
shutil.rmtree(temp)
return net
def refresh_weights(self):
from zoo.util.tf import export_tf
export_tf(self.sess,
self.export_dir,
inputs=self.inputs,
outputs=self.grads + self.outputs)
self.training_helper_layer = TFTrainingHelper(self.export_dir,
self.session_config)
def from_session(sess, inputs, outputs):
temp = tempfile.mkdtemp()
try:
export_tf(sess, temp, inputs, outputs)
net = TFNet.from_export_folder(temp)
finally:
import shutil
shutil.rmtree(temp)
return net
def from_session(sess, inputs, outputs,
generate_backward=False, allow_non_differentiable_input=True):
temp = tempfile.mkdtemp()
try:
export_tf(sess, temp, inputs, outputs,
generate_backward, allow_non_differentiable_input)
net = TFNet.from_export_folder(temp)
finally:
import shutil
shutil.rmtree(temp)
return net
def create(loss, sess, inputs, grads, variables, graph, tensors_with_value,
session_config, metrics, updates):
import tensorflow as tf
from zoo.util.tf import export_tf
inputs, additional_values = TFModel._expand_inputs(
inputs, tensors_with_value, loss)
session_config = TFModel._process_session_config(session_config)
grads = TFModel._process_grads(graph, grads)
outputs, val_methods = TFModel._process_metrics(
graph, metrics, loss, inputs)
assign, variable_placeholders = TFModel._process_variables(
graph, variables)
export_dir = tempfile.mkdtemp()
export_tf(sess, export_dir, inputs=inputs, outputs=grads + outputs)
variable_names = [v.name for v in variables]
grad_names = [g.name for g in grads]
output_names = [o.name for o in outputs]
def to_floats(vs):
return [float(v) for v in vs]
meta = {
"input_names": [i.name for i in inputs],
"output_names": output_names,
"variables": variable_names,
"grad_variables": grad_names,
"default_tensor_values": [to_floats(v) for v in additional_values]
}
with open(os.path.join(export_dir, "training_meta.json"), "w") as f:
f.write(json.dumps(meta))
training_helper_layer = TFTrainingHelper(export_dir, session_config,
assign, variable_placeholders,
sess)
criterion = IdentityCriterion()
return TFModel(training_helper_layer, criterion, val_methods)
def ckpt_to_frozen_graph(options):
with tf.gfile.GFile(options.pbPath, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
var_list_name = [
node.name + ":0" for node in graph_def.node
if node.op in ["Variable", "VariableV2", "VarHandleOp"]
]
# now build the graph in the memory and visualize it
with tf.Session() as sess:
graph = tf.get_default_graph()
tf.import_graph_def(graph_def, name="")
var_list = [graph.get_tensor_by_name(name) for name in var_list_name]
for v in var_list:
tf.add_to_collection(tf.GraphKeys.TRAINABLE_VARIABLES, v)
saver = tf.train.Saver(var_list)
saver.restore(sess, options.ckptPath)
input_names = options.inputsName.split(",")
output_names = options.outputsName.split(",")
input_tensors = [
graph.get_tensor_by_name(name) for name in input_names
]
output_tensors = [
graph.get_tensor_by_name(name) for name in output_names
]
export_tf(sess,
options.outputDir,
inputs=input_tensors,
outputs=output_tensors)
def __init__(self, loss, optim_method, sess=None):
self.optim_method = optim_method
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
grads_vars = tf.train.GradientDescentOptimizer(0).compute_gradients(
loss)
variables = []
grads = []
for (grad, var) in grads_vars:
variables.append(var)
grads.append(grad)
self.export_dir = tempfile.mkdtemp()
all_required_inputs = _find_placeholders([loss])
self.dataset = tf.get_collection(all_required_inputs[0].name)[0]
if self.dataset.batch_size <= 0:
raise ValueError(
"You should set batch_size instead of batch_per_core for training"
)
self.inputs = self.dataset.tensors
_check_the_same(all_required_inputs, self.inputs)
export_tf(self.sess,
self.export_dir,
inputs=self.inputs,
outputs=grads + [loss])
variable_names = [v.name for v in variables]
grad_names = [g.name for g in grads]
meta = {
"input_names": [i.name for i in self.inputs],
"output_names": [loss.name],
"variables": variable_names,
"grad_variables": grad_names
}
with open(os.path.join(self.export_dir, "training_meta.json"),
"w") as f:
f.write(json.dumps(meta))
self.training_helper_layer = TFTrainingHelper(self.export_dir)
self.variable_placeholders = []
assigns = []
for v in variables:
p = tf.placeholder(dtype=tf.float32, shape=v.shape)
a = tf.assign(v, p)
self.variable_placeholders.append(p)
assigns.append(a)
self.assign = tf.group(*assigns)
data = self.dataset.rdd
batch_size = self.dataset.batch_size
sample_rdd = data.map(
lambda t: Sample.from_ndarray(t, [np.array([0.0])]))
self.optimizer = Optimizer.create(self.training_helper_layer,
sample_rdd,
IdentityCriterion(),
batch_size=batch_size,
optim_method=self.optim_method)
def create(loss, sess, inputs, grads, variables, graph, tensors_with_value,
session_config, metrics):
import tensorflow as tf
from zoo.util.tf import export_tf
additional_inputs = []
additional_values = []
all_required_inputs = _find_placeholders([loss])
all_required_inputs_names = [v.name for v in all_required_inputs]
if tensors_with_value:
for t, v in tensors_with_value.items():
if t.name in all_required_inputs_names:
additional_inputs.append(t)
additional_values.append(v)
if not isinstance(inputs, list):
inputs = nest.flatten(inputs)
inputs = inputs + additional_inputs
if session_config is not None:
import tensorflow as tf
assert isinstance(session_config, tf.ConfigProto),\
"session_config should be a tf.ConfigProto"
session_config.use_per_session_threads = True
session_config = session_config
from zoo.util.tf import process_grad
grads = [process_grad(grad) for grad in grads]
outputs = []
val_methods = None
if metrics is not None:
idx = 0
val_methods = []
for metric_name in metrics:
metric = metrics[metric_name]
if tf.is_numeric_tensor(metric):
outputs.append(metric)
val_methods.append(StatelessMetric(metric_name, idx))
idx += 1
else:
outputs += metric.outputs
with graph.as_default():
val_labels = [tf.identity(v) for v in metric.labels]
outputs += val_labels
method = TFValidationMethod(
metric.val_method, metric_name,
list(range(idx, idx + len(metric.outputs))),
list(
range(idx + len(metric.outputs), idx +
len(metric.outputs) + len(val_labels))))
val_methods.append(method)
idx += len(metric.outputs) + len(val_labels)
with graph.as_default():
real_batch_size = tf.shape(inputs[0])[0]
outputs.append(real_batch_size)
outputs.append(loss)
export_dir = tempfile.mkdtemp()
export_tf(sess, export_dir, inputs=inputs, outputs=grads + outputs)
variable_names = [v.name for v in variables]
grad_names = [g.name for g in grads]
output_names = [o.name for o in outputs]
def to_floats(vs):
return [float(v) for v in vs]
meta = {
"input_names": [i.name for i in inputs],
"output_names": output_names,
"variables": variable_names,
"grad_variables": grad_names,
"default_tensor_values": [to_floats(v) for v in additional_values]
}
with open(os.path.join(export_dir, "training_meta.json"), "w") as f:
f.write(json.dumps(meta))
variable_placeholders = []
with graph.as_default():
assigns = []
for v in variables:
p = tf.placeholder(dtype=tf.float32, shape=v.shape)
a = tf.assign(v, p)
variable_placeholders.append(p)
assigns.append(a)
assign = tf.group(*assigns)
assign = assign
training_helper_layer = TFTrainingHelper(export_dir, session_config,
assign, variable_placeholders)
criterion = IdentityCriterion()
return TFModel(training_helper_layer, criterion, val_methods)
def __init__(self,
loss,
optim_method,
sess=None,
dataset=None,
inputs=None,
grads=None,
variables=None,
graph=None,
val_outputs=None,
val_labels=None,
val_method=None,
add_sample_weights_num=0):
import tensorflow as tf
from zoo.util.tf import export_tf
'''
TFOptimizer is used for distributed training of tensorflow
on Spark/BigDL.
:param loss: The loss tensor of the tensorflow model, should be a scalar
:param optim_method: the optimization method to be used, such as bigdl.optim.optimizer.Adam
:param sess: the current tensorflow Session, if you want to used a pre-trained model, you
should use the Session to load the pre-trained variables and pass it to TFOptimizer.
'''
if dataset is None:
args = TFOptimizer._get_arguments_from_loss(
loss, optim_method, sess, val_outputs, val_labels, val_method)
loss, optim_method, sess, dataset, inputs = args[:5]
grads, variables, graph, val_outputs, val_labels, val_method = args[
5:]
self.optim_method = optim_method
self.sess = sess
self.dataset = dataset
self.inputs = inputs
self.graph = graph
if self.dataset.batch_size <= 0:
raise ValueError(
"You should set batch_size instead of batch_per_thread for training"
)
if val_outputs is not None and val_labels is not None:
with self.graph.as_default():
val_labels = [tf.identity(v) for v in val_labels]
outputs = val_outputs + val_labels + [loss]
else:
outputs = [loss]
self.export_dir = tempfile.mkdtemp()
export_tf(self.sess,
self.export_dir,
inputs=self.inputs,
outputs=grads + outputs)
variable_names = [v.name for v in variables]
grad_names = [g.name for g in grads]
output_names = [o.name for o in outputs]
meta = {
"input_names": [i.name for i in self.inputs],
"output_names": output_names,
"variables": variable_names,
"grad_variables": grad_names
}
with open(os.path.join(self.export_dir, "training_meta.json"),
"w") as f:
f.write(json.dumps(meta))
self.variable_placeholders = []
with self.graph.as_default():
assigns = []
for v in variables:
p = tf.placeholder(dtype=tf.float32, shape=v.shape)
a = tf.assign(v, p)
self.variable_placeholders.append(p)
assigns.append(a)
assign = tf.group(*assigns)
self.assign = assign
self.training_helper_layer = TFTrainingHelper(self.export_dir)
data = self.dataset.rdd
batch_size = self.dataset.batch_size
sample_rdd = data.map(lambda t: Sample.from_ndarray(
t + [np.array(1.0)] * add_sample_weights_num, [np.array([0.0])]))
self.optimizer = Optimizer.create(self.training_helper_layer,
sample_rdd,
IdentityCriterion(),
batch_size=batch_size,
optim_method=self.optim_method)
if val_outputs is not None and val_labels is not None:
val_sample_rdd = self.dataset.val_rdd\
.map(lambda t: Sample.from_ndarray(t + [np.array(1.0)] * add_sample_weights_num,
[np.array([0.0])]))
val_method = [
TFValidationMethod(m, len(val_outputs), len(val_labels))
for m in to_list(val_method)
]
self.optimizer.set_validation(self.dataset.batch_size,
val_sample_rdd, EveryEpoch(),
val_method)