def _get_arguments_from_loss(loss, optim_method, session, val_outputs, val_labels, val_method): import tensorflow as tf if session is None: sess = tf.Session() sess.run(tf.global_variables_initializer()) else: sess = session grads_vars = tf.train.GradientDescentOptimizer(0).compute_gradients( loss) grads_vars.sort(key=lambda grad_var: grad_var[1].name) variables = [] grads = [] for (grad, var) in grads_vars: if grad is not None: variables.append(var) grads.append(grad) all_required_inputs = _find_placeholders([loss]) dataset = tf.get_collection(all_required_inputs[0].name)[0] inputs = nest.flatten(dataset._original_tensors) return [ loss, optim_method, sess, dataset, inputs, grads, variables, loss.graph, val_outputs, val_labels, val_method ]
def _get_datasets_and_inputs(outputs): import tensorflow as tf all_required_inputs = _find_placeholders(outputs) dataset = tf.get_collection(all_required_inputs[0].name)[0] inputs = dataset.tensors _check_the_same(all_required_inputs, inputs) return dataset, inputs
def _expand_inputs(inputs, tensors_with_value, loss): 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 return inputs, additional_values
def _get_arguments_from_loss(loss, optim_method, session, val_outputs, val_labels, val_method): import tensorflow as tf if session is None: sess = tf.Session() sess.run(tf.global_variables_initializer()) else: sess = session grads, variables = TFOptimizer._get_vars_grads(loss) all_required_inputs = _find_placeholders([loss]) dataset = tf.get_collection(all_required_inputs[0].name)[0] inputs = nest.flatten(dataset._original_tensors) return [ loss, optim_method, sess, dataset, inputs, grads, variables, loss.graph, val_outputs, val_labels, val_method ]
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, val_split=0.0, tensors_with_value=None, session_config=None, clip_norm=None, clip_value=None): ''' 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. ''' import tensorflow as tf from zoo.util.tf import export_tf 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:] 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) self.optim_method = optim_method self.sess = sess self.dataset = dataset self.inputs = inputs + additional_inputs self.graph = graph 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 self.session_config = session_config self.clip_norm = clip_norm if clip_value is not None and not isinstance(clip_value, tuple): raise ValueError("The clip_value argument should be a tuple (min_value, max_value)") self.clip_constant = clip_value from zoo.util.tf import process_grad grads = [process_grad(grad) for grad in grads] 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.grads = grads self.outputs = outputs self.export_dir = tempfile.mkdtemp() export_tf(self.sess, self.export_dir, inputs=self.inputs, outputs=self.grads + self.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 self.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(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 try: self.training_helper_layer = TFTrainingHelper(self.export_dir, session_config) except Py4JJavaError as e: if "expects to be colocated with unknown node" in str(e): raise Exception(""" If you are using the embedding layer in tf.keras, then this is a known issue of TensorFlow, see https://github.com/tensorflow/tensorflow/issues/21889. Please add zoo.util.tf.variable_creator_scope before model construction. For example: from zoo.util.tf import variable_creator_scope with variable_creator_scope(): model = tf.keras.models.Sequential([ tf.keras.layers.Embedding(1, 1, input_length=1)]) """) else: raise e batch_size = self.dataset.batch_size sample_rdd = self.dataset.get_training_data() if val_outputs is not None and val_labels is not None: val_rdd = self.dataset.get_validation_data() if val_rdd is not None: val_method = [TFValidationMethod(m, len(val_outputs), len(val_labels)) for m in to_list(val_method)] training_rdd = sample_rdd elif val_split != 0.0: training_rdd, val_rdd = sample_rdd.randomSplit([1 - val_split, val_split]) val_method = [TFValidationMethod(m, len(val_outputs), len(val_labels)) for m in to_list(val_method)] else: raise ValueError("Validation data is not specified. Please set " + "val rdd in TFDataset, or set val_split larger than zero") self.optimizer = Optimizer.create(self.training_helper_layer, training_rdd, IdentityCriterion(), batch_size=batch_size, optim_method=self.optim_method) self.optimizer.set_validation(self.dataset.batch_size, val_rdd, EveryEpoch(), val_method) else: training_rdd = sample_rdd self.optimizer = Optimizer.create(self.training_helper_layer, training_rdd, IdentityCriterion(), batch_size=batch_size, optim_method=self.optim_method) if self.clip_norm: self.optimizer.set_gradclip_l2norm(self.clip_norm) if self.clip_constant: min_value, max_value = self.clip_constant self.optimizer.set_gradclip_const(min_value, max_value)
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)