def run_tf_model(args, dense_variables, vocabulary_tensors, samples, labels): model = TfDenseDemo(global_batch_size=args.global_batch_size, vocabulary_size=args.max_vocabulary_size_per_gpu * hvd.size(), slot_num=args.slot_num, nnz_per_slot=args.nnz_per_slot, num_dense_layers=args.num_dense_layers, num_dense_units=args.num_dense_units, embedding_vec_size=args.embedding_vec_size) #model.build(input_shape=(args.global_batch_size, args.slot_num * args.nnz_per_slot * args.embedding_vec_size)) model(samples, training=False) for i in range(args.num_dense_layers): model.dense_layers[i].trainable_variables[0].assign( dense_variables[0][i]) model.dense_layers[i].trainable_variables[1].assign( dense_variables[1][i]) vocabulary_table = tf.concat(vocabulary_tensors, axis=0) for i in range(hvd.size()): model.params.assign(vocabulary_table) optimizer = utils.get_dense_optimizer(args.optimizer)(learning_rate=0.1) if args.mixed_precision: optimizer = tf.keras.mixed_precision.LossScaleOptimizer( optimizer, initial_scale=1024) loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True) @tf.function def _train_step(inputs, labels): with tf.GradientTape() as tape: logit = model(inputs, training=True) loss = loss_fn(labels, logit) if args.mixed_precision: _loss = optimizer.get_scaled_loss(loss) else: _loss = loss grads = tape.gradient(_loss, model.trainable_variables) if args.mixed_precision: grads = optimizer.get_unscaled_gradients(grads) optimizer.apply_gradients(zip(grads, model.trainable_variables)) return loss loss_list = [] for i in range(args.iter_num): loss = _train_step(samples, labels) loss_list.append(loss) print("[INFO]: Iteration: {}, loss={}".format(i, loss)) return loss_list
def test_tf_dense_model(args, init_tensors, *random_samples): dataset = utils.tf_dataset(*random_samples, batchsize=args.global_batch_size, to_sparse_tensor=False, repeat=1) loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True) tf_dense_demo = TfDenseDemo(init_tensors, args.global_batch_size, args.slot_num, args.nnz_per_slot, args.embedding_vec_size) optimizer = utils.get_dense_optimizer(args.optimizer)(learning_rate=0.1) if args.mixed_precision: optimizer = tf.keras.mixed_precision.LossScaleOptimizer( optimizer, initial_scale=1024) @tf.function def _train_step(inputs, labels): with tf.GradientTape() as tape: logit, embedding_vector = tf_dense_demo(inputs, training=True) loss = loss_fn(labels, logit) if args.mixed_precision: _loss = optimizer.get_scaled_loss(loss) else: _loss = loss grads = tape.gradient(_loss, tf_dense_demo.trainable_variables) if args.mixed_precision: grads = optimizer.get_unscaled_gradients(grads) optimizer.apply_gradients(zip(grads, tf_dense_demo.trainable_variables)) return loss, embedding_vector tf_results = list() for i, (input_tensors, labels) in enumerate(dataset): print("-" * 30, str(i), "-" * 30) loss, embedding_vector = _train_step(input_tensors, labels) print("[INFO]: iteration {}, loss {}".format(i, loss)) tf_results.append(embedding_vector.numpy()) if not hasattr(args, "task_id"): args.task_id = 0 if 1 == args.save_params and args.task_id == 0: filepath = r"./embedding_variables/" utils.save_to_file(os.path.join(filepath, r"tf_variable.file"), tf_dense_demo.params.numpy()) return tf_results
def test_sok_dense_demo(args, init_tensors, *random_samples): port = 12345 os.environ["TF_CONFIG"] = json.dumps({ "cluster": {"worker": [args.ips[i] + ":" + str(port + i) for i in range(args.worker_num)]}, "task": {"type": "worker", "index": args.task_id} }) strategy = tf.distribute.MultiWorkerMirroredStrategy() with strategy.scope(): sok.Init(global_batch_size=args.global_batch_size) sok_dense_demo = SOKDenseDemo(max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size=args.embedding_vec_size, slot_num=args.slot_num, nnz_per_slot=args.nnz_per_slot, use_hashtable=args.use_hashtable) emb_opt = utils.get_embedding_optimizer(args.optimizer)(learning_rate=0.1) dense_opt = utils.get_dense_optimizer(args.optimizer)(learning_rate=0.1) if args.mixed_precision: emb_opt = tf.keras.mixed_precision.LossScaleOptimizer(emb_opt, initial_scale=1024) sok_saver = sok.Saver() if 1 == args.restore_params: filepath = r"./embedding_variables" sok_saver.restore_from_file(sok_dense_demo.embedding_layer.embedding_variable, filepath) else: sok_saver.load_embedding_values(sok_dense_demo.embedding_layer.embedding_variable, init_tensors) loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss(loss, global_batch_size=args.global_batch_size) return tf.cast(loss, _dtype) @tf.function def _train_step(inputs, labels): with tf.GradientTape() as tape: logit, embedding_vector = sok_dense_demo(inputs, training=True) loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = emb_opt.get_scaled_loss(loss) else: _loss = loss embedding_variables, other_variable = sok.split_embedding_variable_from_others(sok_dense_demo.trainable_variables) grads, emb_grads = tape.gradient(_loss, [other_variable, embedding_variables]) if args.mixed_precision: grads = emb_opt.get_unscaled_gradients(grads) emb_grads = emb_opt.get_unscaled_gradients(emb_grads) if "plugin" not in args.optimizer: with sok.OptimizerScope(embedding_variables): emb_opt.apply_gradients(zip(emb_grads, embedding_variables), experimental_aggregate_gradients=False) else: emb_opt.apply_gradients(zip(emb_grads, embedding_variables), experimental_aggregate_gradients=False) dense_opt.apply_gradients(zip(grads, other_variable)) return loss, embedding_vector sok_results = list() def _dataset_fn(input_context): replica_batch_size = input_context.get_per_replica_batch_size(args.global_batch_size) dataset = utils.tf_dataset(*random_samples, batchsize=replica_batch_size, to_sparse_tensor=False, repeat=1) return dataset dataset = strategy.distribute_datasets_from_function(_dataset_fn) for i, (input_tensors, replica_labels) in enumerate(dataset): print("-"*30, "step ", str(i), "-"*30) loss, embedding_vector = strategy.run(_train_step, args=(input_tensors, replica_labels)) loss = strategy.reduce("sum", loss, axis=None) print("[INFO]: iteration {}, loss {}".format(i, loss)) sok_results.append(embedding_vector) # save params to file. if 1 == args.save_params: filepath = r"./embedding_variables" utils.try_make_dirs(filepath, chief=(True if args.task_id == 0 else False)) sok_saver.dump_to_file(sok_dense_demo.embedding_layer.embedding_variable, filepath) return sok_results, sok_dense_demo.embedding_layer.embedding_variable.values[0].m_var_name
def run_sok_model(args, dense_variables, vocabulary_tensors, samples, labels): # split sample and labels assert (args.global_batch_size % hvd.size() == 0) local_batch_size = args.global_batch_size // hvd.size() local_id = hvd.local_rank() samples = samples[local_id * local_batch_size:(local_id + 1) * local_batch_size] labels = labels[local_id * local_batch_size:(local_id + 1) * local_batch_size] sok.Init(global_batch_size=args.global_batch_size) model = SOKDenseDemo( max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size=args.embedding_vec_size, slot_num=args.slot_num, nnz_per_slot=args.nnz_per_slot, num_dense_layers=args.num_dense_layers, num_dense_units=args.num_dense_units) #model.build(input_shape=(local_batch_size, args.slot_num * args.nnz_per_slot * args.embedding_vec_size)) model(samples, training=False) for i in range(args.num_dense_layers): model.dense_layers[i].trainable_variables[0].assign( dense_variables[0][i]) model.dense_layers[i].trainable_variables[1].assign( dense_variables[1][i]) sok_saver = sok.Saver() init_tensors = [tensor.numpy() for tensor in vocabulary_tensors] sok_saver.load_embedding_values(model.embedding_layer.embedding_variable, init_tensors) embedding_optimizer = utils.get_embedding_optimizer( args.optimizer)(learning_rate=0.1) dense_optimizer = utils.get_dense_optimizer( args.optimizer)(learning_rate=0.1) if args.mixed_precision: embedding_optimizer = tf.keras.mixed_precision.LossScaleOptimizer( embedding_optimizer, initial_scale=1024) loss_fn = tf.keras.losses.BinaryCrossentropy( from_logits=True, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss( loss, global_batch_size=args.global_batch_size) return tf.cast(loss, dtype=_dtype) @tf.function def _train_step(inputs, labels, first_batch): with tf.GradientTape() as tape, tf.GradientTape() as emb_tape: logit = model(inputs, training=True) replica_loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = embedding_optimizer.get_scaled_loss(replica_loss) else: _loss = replica_loss tape = hvd.DistributedGradientTape(tape) emb_variable, other_variable = sok.split_embedding_variable_from_others( model.trainable_variables) emb_grads = emb_tape.gradient(_loss, emb_variable) grads = tape.gradient(_loss, other_variable) if args.mixed_precision: emb_grads = embedding_optimizer.get_unscaled_gradients(emb_grads) grads = embedding_optimizer.get_unscaled_gradients(grads) if 'plugin' not in args.optimizer: with sok.OptimizerScope(emb_variable): embedding_optimizer.apply_gradients( zip(emb_grads, emb_variable), experimental_aggregate_gradients=False) else: embedding_optimizer.apply_gradients( zip(emb_grads, emb_variable), experimental_aggregate_gradients=False) dense_optimizer.apply_gradients(zip(grads, other_variable)) # Note: broadcast should be done after the first gradient step to ensure optimizer initialization. if first_batch: hvd.broadcast_variables(other_variable, root_rank=0) hvd.broadcast_variables(dense_optimizer.variables(), root_rank=0) return replica_loss loss_list = [] for i in range(args.iter_num): loss = _train_step(samples, labels, i == 0) loss_list.append(loss) print("[INFO]: Iteration: {}, loss={}".format(i, loss)) return loss_list
def test_sok_multi_dense_emb(args): comm_options = tf.distribute.experimental.CommunicationOptions( bytes_per_pack=0, timeout_seconds=None, implementation=tf.distribute.experimental.CommunicationImplementation. NCCL) if args.worker_num == 1: strategy = tf.distribute.MirroredStrategy() else: port = 12345 os.environ["TF_CONFIG"] = json.dumps({ "cluster": { "worker": [ "localhost" + ":" + str(port + i) for i in range(args.worker_num) ] }, "task": { "type": "worker", "index": args.task_id } }) strategy = tf.distribute.MultiWorkerMirroredStrategy( communication_options=comm_options) replica_batch_size = args.global_batch_size // (args.worker_num * 1) dataset = utility.TFDataset(filename=args.file_prefix + str(args.task_id) + ".file", batchsize=replica_batch_size, as_sparse_tensor=False, repeat=1) dataset = dataset.prefetch(tf.data.AUTOTUNE) dynamic_input = True if args.dynamic_input == 1 else False with strategy.scope(): sok.Init(global_batch_size=args.global_batch_size) model = SOKDenseModel( max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size_list=args.embedding_vec_size_list, slot_num_list=args.slot_num_list, nnz_per_slot_list=[ args.nnz_per_slot for _ in range(len(args.slot_num_list)) ], num_dense_layers=args.num_dense_layers, dynamic_input=dynamic_input) emb_opt = utils.get_embedding_optimizer( args.optimizer)(learning_rate=0.1) dense_opt = utils.get_dense_optimizer( args.optimizer)(learning_rate=0.1) if args.mixed_precision: emb_opt = tf.keras.mixed_precision.LossScaleOptimizer( emb_opt, initial_scale=1024) # set initial value to embedding variables. sok_saver = sok.Saver() for i, layer in enumerate(model.embedding_layers): init_tensors = utils.get_ones_tensor( max_vocab_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size=args.embedding_vec_size_list[i], num=args.worker_num) sok_saver.load_embedding_values(layer.embedding_variable, init_tensors) loss_fn = tf.keras.losses.BinaryCrossentropy( from_logits=True, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss( loss, global_batch_size=args.global_batch_size) return tf.cast(loss, _dtype) @tf.function def _train_step(inputs, labels): with tf.GradientTape() as tape: logit, all_vectors = model(inputs, training=True) loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = emb_opt.get_scaled_loss(loss) else: _loss = loss emb_variable, other_variable = sok.split_embedding_variable_from_others( model.trainable_variables) grads, emb_grads = tape.gradient(_loss, [other_variable, emb_variable]) if args.mixed_precision: grads = emb_opt.get_unscaled_gradients(grads) emb_grads = emb_opt.get_unscaled_gradients(emb_grads) if "plugin" not in args.optimizer: with sok.OptimizerScope(emb_variable): emb_opt.apply_gradients(zip(emb_grads, emb_variable), experimental_aggregate_gradients=False) else: emb_opt.apply_gradients(zip(emb_grads, emb_variable), experimental_aggregate_gradients=False) with tf.control_dependencies(emb_grads): # mannually all-reduce dense gradients replica_context = tf.distribute.get_replica_context() grads = replica_context.all_reduce("sum", grads, options=comm_options) dense_opt.apply_gradients(zip(grads, other_variable), experimental_aggregate_gradients=False) # manually all-reduce loss, it is ok, because replica_loss has already been used to # update local variables. loss = replica_context.all_reduce(tf.distribute.ReduceOp.SUM, loss, options=comm_options) return loss, all_vectors, logit # save its results sok_results = list() for i, (inputs, labels) in enumerate(dataset): if args.stop_iter >= 0 and i >= args.stop_iter: break total_loss, all_vectors, logit = strategy.run(_train_step, args=(inputs, labels)) print("[INFO]: Iteration: {}, loss={}".format(i, total_loss)) with tf.device("CPU:0"): sok_results.append(all_vectors) return sok_results
def test_sok_demo(args, init_tensors, *random_samples): strategy = tf.distribute.MirroredStrategy() with strategy.scope(): result = sok.Init(global_batch_size=args.global_batch_size) embedding_initializer = tf.keras.initializers.Ones( ) if args.use_tf_initializer else None plugin_demo = SOKDemo( combiner=args.combiner, max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, slot_num=args.slot_num, max_nnz=args.max_nnz, embedding_vec_size=args.embedding_vec_size, use_hashtable=args.use_hashtable, key_dtype=args.key_dtype, embedding_initializer=embedding_initializer) emb_opt = utils.get_embedding_optimizer( args.optimizer)(learning_rate=0.1) dense_opt = utils.get_dense_optimizer( args.optimizer)(learning_rate=0.1) if args.mixed_precision: emb_opt = tf.keras.mixed_precision.LossScaleOptimizer( emb_opt, initial_scale=1024) plugin_saver = sok.Saver() if (1 == args.restore_params): # restore from trained parameters filepath = r"./embedding_variables" plugin_saver.restore_from_file( plugin_demo.embedding_layer.embedding_variable, filepath) else: # initialize using randomized initial value if not args.use_tf_initializer and init_tensors: status = plugin_saver.load_embedding_values( plugin_demo.embedding_layer.embedding_variable, init_tensors) loss_fn = tf.keras.losses.BinaryCrossentropy( from_logits=True, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss( loss, global_batch_size=args.global_batch_size) return tf.cast(loss, _dtype) @tf.function def _train_step(inputs, labels): with tf.GradientTape() as tape: logit, embedding_vector = plugin_demo(inputs, training=True) loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = emb_opt.get_scaled_loss(loss) else: _loss = loss embedding_variables, other_variable = sok.split_embedding_variable_from_others( plugin_demo.trainable_variables) grads, emb_grads = tape.gradient(_loss, [other_variable, embedding_variables]) if args.mixed_precision: grads = emb_opt.get_unscaled_gradients(grads) emb_grads = emb_opt.get_unscaled_gradients(emb_grads) with tf.control_dependencies([*emb_grads]): # in case NCCL runs concurrently via SOK and TF if 'plugin' not in args.optimizer: with sok.OptimizerScope(embedding_variables): emb_opt.apply_gradients( zip(emb_grads, embedding_variables), experimental_aggregate_gradients=False) else: emb_opt.apply_gradients(zip(emb_grads, embedding_variables), experimental_aggregate_gradients=False) dense_opt.apply_gradients(zip(grads, other_variable)) return loss, embedding_vector sok_results = list() def _dataset_fn(input_context): replica_batch_size = input_context.get_per_replica_batch_size( args.global_batch_size) dataset = utils.tf_dataset(*random_samples, batchsize=replica_batch_size, to_sparse_tensor=True, repeat=1, args=args) dataset = dataset.shard(input_context.num_input_pipelines, input_context.input_pipeline_id) return dataset dataset = strategy.distribute_datasets_from_function(_dataset_fn) for i, (sparse_tensors, replica_labels) in enumerate(dataset): print("-" * 30, "step ", str(i), "-" * 30) loss, embedding_vector = strategy.run(_train_step, args=(sparse_tensors, replica_labels)) loss = strategy.reduce("sum", loss, axis=None) print("[INFO]: iteration {}, loss {}".format(i, loss)) sok_results.append(embedding_vector) # save params to file. if 1 == args.save_params: filepath = r"./embedding_variables/" utils.try_make_dirs(filepath) plugin_saver.dump_to_file( plugin_demo.embedding_layer.embedding_variable, filepath) return sok_results, plugin_demo.embedding_layer.embedding_variable.values[ 0].m_var_name
def get_sok_results(args, init_tensors, *random_samples): if args.distributed_tool == "onedevice": strategy = strategy_wrapper.OneDeviceStrategy() elif args.distributed_tool == "horovod": import horovod.tensorflow as hvd hvd.init() strategy = strategy_wrapper.HorovodStrategy() else: raise ValueError(f"{args.distributed_tool} is not supported.") with strategy.scope(): sok_init_op = sok.Init(global_batch_size=args.global_batch_size) embedding_initializer = tf.keras.initializers.Ones( ) if args.use_tf_initializer else None sok_dense_demo = SOKDemo( max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size=args.embedding_vec_size, slot_num=args.slot_num, nnz_per_slot=args.nnz_per_slot, use_hashtable=args.use_hashtable, dynamic_input=args.dynamic_input, num_of_dense_layers=0, key_dtype=args.key_dtype, embedding_initializer=embedding_initializer) emb_opt = utils.get_embedding_optimizer( args.optimizer)(learning_rate=0.1) dense_opt = utils.get_dense_optimizer( args.optimizer)(learning_rate=0.1) if args.mixed_precision: emb_opt = sok.tf.keras.mixed_precision.LossScaleOptimizer( emb_opt, 1024) sok_saver = sok.Saver() restore_op = list() for i, embedding_layer in enumerate(sok_dense_demo.embedding_layers): control_inputs = [restore_op[-1]] if restore_op else None with tf.control_dependencies(control_inputs): if args.restore_params: filepath = r"./embedding_variables" op = sok_saver.restore_from_file( embedding_layer.embedding_variable, filepath) else: if not args.use_tf_initializer: op = sok_saver.load_embedding_values( embedding_layer.embedding_variable, init_tensors[i]) else: op = tf.constant(1.0) restore_op.append(op) loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction='none') def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss( loss, global_batch_size=args.global_batch_size) return tf.cast(loss, _dtype) def _train_step(inputs, labels, training): def _step_fn(inputs, labels): logit, embedding_vector = sok_dense_demo(inputs, training=training) loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = emb_opt.get_scaled_loss(loss) else: _loss = loss emb_var, other_var = sok.split_embedding_variable_from_others( sok_dense_demo.trainable_variables) grads = tf.gradients( _loss, emb_var + other_var, colocate_gradients_with_ops=True, unconnected_gradients=tf.UnconnectedGradients.NONE) emb_grads, other_grads = grads[:len(emb_var)], grads[len(emb_var):] if args.mixed_precision: other_grads = emb_opt.get_unscaled_gradients(other_grads) emb_grads = emb_opt.get_unscaled_gradients(emb_grads) if "plugin" in args.optimizer: emb_train_op = emb_opt.apply_gradients(zip(emb_grads, emb_var)) else: with sok.OptimizerScope(emb_var): emb_train_op = emb_opt.apply_gradients( zip(emb_grads, emb_var)) with tf.control_dependencies([*emb_grads]): # in case NCCL runs concurrently via SOK and horovod other_grads = strategy.reduce("sum", other_grads) other_train_op = dense_opt.apply_gradients( zip(other_grads, other_var)) with tf.control_dependencies([emb_train_op, other_train_op]): total_loss = strategy.reduce("sum", loss) total_loss = tf.identity(total_loss) return total_loss, embedding_vector return strategy.run(_step_fn, inputs, labels) replica_batch_size = args.global_batch_size // args.gpu_num dataset = utils.tf_dataset(*random_samples, batchsize=replica_batch_size, to_sparse_tensor=False, repeat=1, args=args) train_iterator = dataset.make_initializable_iterator() iterator_init = train_iterator.initializer inputs, labels = train_iterator.get_next() graph_results = _train_step(inputs, labels, training=True) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) if "plugin" in args.optimizer: init_op = tf.group(init_op, emb_opt.initializer) save_op = list() for i, embedding_layer in enumerate(sok_dense_demo.embedding_layers): control_inputs = [save_op[-1]] if save_op else None with tf.control_dependencies(control_inputs): if args.save_params: filepath = r"./embedding_variables/" utils.try_make_dirs(filepath) op = sok_saver.dump_to_file(embedding_layer.embedding_variable, filepath) else: op = tf.constant(1.0) save_op.append(op) sok_results = list() config = tf.ConfigProto() config.log_device_placement = False with tf.Session(config=config) as sess: sess.run(sok_init_op) sess.run([init_op, iterator_init]) sess.run(restore_op) sess.graph.finalize() for step in range(args.iter_num): loss_v, emb_vector_v = sess.run([*graph_results]) print("*" * 80) print(f"Step: {step}, loss: {loss_v}" ) #", embedding_vector:\n{emb_vector_v}") sok_results.append(emb_vector_v) sess.run(save_op) name = list() for embedding_layer in sok_dense_demo.embedding_layers: name.append(embedding_layer.embedding_variable.m_var_name) return sok_results, name
def get_tf_results(args, init_tensors, *random_samples): graph = tf.Graph() with graph.as_default(): tf_dense_demo = TFDemo( vocabulary_size=args.max_vocabulary_size_per_gpu * args.gpu_num, slot_num=args.slot_num, nnz_per_slot=args.nnz_per_slot, embedding_vec_size=args.embedding_vec_size, num_of_dense_layers=0, use_hashtable=False, dynamic_input=False) optimizer = utils.get_dense_optimizer( args.optimizer)(learning_rate=0.1) if args.mixed_precision: optimizer = sok.tf.keras.mixed_precision.LossScaleOptimizer( optimizer, 1024) loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True) def _train_step(inputs, labels, training): logit, embedding_vector = tf_dense_demo(inputs, training=training) loss = loss_fn(labels, logit) if args.mixed_precision: _loss = optimizer.get_scaled_loss(loss) else: _loss = loss grads = tf.gradients( _loss, tf_dense_demo.trainable_variables, colocate_gradients_with_ops=True, unconnected_gradients=tf.UnconnectedGradients.NONE) if args.mixed_precision: grads = optimizer.get_unscaled_gradients(grads) train_op = optimizer.apply_gradients( zip(grads, tf_dense_demo.trainable_variables)) with tf.control_dependencies([train_op]): loss = tf.identity(loss) return loss, embedding_vector dataset = utils.tf_dataset(*random_samples, batchsize=args.global_batch_size, to_sparse_tensor=False, repeat=1) train_iterator = dataset.make_initializable_iterator() iterator_init = train_iterator.initializer inputs, labels = train_iterator.get_next() graph_results = _train_step(inputs, labels, training=True) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restore_op = list() for i, embedding_layer in enumerate(tf_dense_demo.embedding_layers): restore_op.append( embedding_layer.embeddings.assign( tf.concat(init_tensors[i], axis=0))) emb_values = list() for embedding_layer in tf_dense_demo.embedding_layers: if args.save_params: filepath = r"./embedding_variables/" utils.try_make_dirs(filepath) emb_values.append(embedding_layer.embeddings.read_value()) else: emb_values = tf.constant(1.0) tf_results = list() with tf.Session(graph=graph) as sess: sess.run([init_op, iterator_init]) sess.run(restore_op) sess.graph.finalize() for step in range(args.iter_num): loss_v, embedding_vector_v = sess.run([*graph_results]) print("*" * 80) print(f"step: {step}, loss: {loss_v}" ) #", embedding_vector:\n{embedding_vector_v}") tf_results.append(embedding_vector_v) emb_values_v = sess.run(emb_values) if args.save_params: for i, value in enumerate(emb_values_v): utils.save_to_file( os.path.join(filepath, r"tf_variable_" + str(i) + r".file"), value) name = list() for embedding_layer in tf_dense_demo.embedding_layers: name.append(embedding_layer.embeddings.name) return tf_results, name
def test_sok_multi_dense_emb(args): assert (args.global_batch_size % args.worker_num == 0) replica_batch_size = args.global_batch_size // (args.worker_num) dataset = utility.TFDataset(filename=args.file_prefix + str(args.task_id) + ".file", batchsize=replica_batch_size, as_sparse_tensor=False, repeat=1) dataset = dataset.prefetch(tf.data.AUTOTUNE) dynamic_input = True if args.dynamic_input == 1 else False # SOK initialize sok.Init(global_batch_size=args.global_batch_size) model = SOKDenseModel( max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size_list=args.embedding_vec_size_list, slot_num_list=args.slot_num_list, nnz_per_slot_list=[ args.nnz_per_slot for _ in range(len(args.slot_num_list)) ], num_dense_layers=args.num_dense_layers, dynamic_input=dynamic_input, use_hashtable=args.use_hashtable) emb_opt = utils.get_embedding_optimizer(args.optimizer)(learning_rate=0.1) dense_opt = utils.get_dense_optimizer(args.optimizer)(learning_rate=0.1) if args.mixed_precision: emb_opt = tf.keras.mixed_precision.LossScaleOptimizer( emb_opt, initial_scale=1024) sok_saver = sok.Saver() for i, layer in enumerate(model.embedding_layers): init_tensors = utils.get_ones_tensor( max_vocab_size_per_gpu=args.max_vocabulary_size_per_gpu, embedding_vec_size=args.embedding_vec_size_list[i], num=args.worker_num) sok_saver.load_embedding_values(layer.embedding_variable, init_tensors) loss_fn = tf.keras.losses.BinaryCrossentropy( from_logits=True, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss = loss_fn(labels, logits) _dtype = loss.dtype loss = tf.cast(loss, tf.float32) loss = tf.nn.compute_average_loss( loss, global_batch_size=args.global_batch_size) return tf.cast(loss, _dtype) @tf.function def _train_step(inputs, labels, first_batch): with tf.GradientTape() as tape: logit, all_vectors = model(inputs, training=True) replica_loss = _replica_loss(labels, logit) if args.mixed_precision: _loss = emb_opt.get_scaled_loss(replica_loss) else: _loss = replica_loss emb_var, other_var = sok.split_embedding_variable_from_others( model.trainable_variables) emb_grads, grads = tape.gradient(_loss, [emb_var, other_var]) if args.mixed_precision: emb_grads = emb_opt.get_unscaled_gradients(emb_grads) grads = emb_opt.get_unscaled_gradients(grads) if "plugin" not in args.optimizer: with sok.OptimizerScope(emb_var): emb_opt.apply_gradients(zip(emb_grads, emb_var), experimental_aggregate_gradients=False) else: emb_opt.apply_gradients(zip(emb_grads, emb_var), experimental_aggregate_gradients=False) with tf.control_dependencies(emb_grads): grads = [hvd.allreduce(grad) for grad in grads] dense_opt.apply_gradients(zip(grads, other_var)) if first_batch: hvd.broadcast_variables(other_var, root_rank=0) hvd.broadcast_variables(dense_opt.variables(), root_rank=0) total_loss = hvd.allreduce(replica_loss) return total_loss, all_vectors sok_results = list() for i, (inputs, labels) in enumerate(dataset): if args.stop_iter >= 0 and i >= args.stop_iter: break total_loss, all_vectors = _train_step(inputs, labels, 0 == i) print("[INFO]: Iteration: {}, loss={}".format(i, total_loss)) sok_results.append(all_vectors) return sok_results