[Fix] Remove TF_DETERMINISTIC_OPS

deepray/__init__.py

@@ -36,16 +36,23 @@
 from deepray.utils import types
 from deepray.utils.ensure_tf_install import _check_tf_version
 from deepray.utils.flags import common_flags
+from deepray.utils.keras_utils import set_random_seed
 from deepray.version import __version__
 
 # _check_tf_version()
 
 logger = logging_util.get_logger()
 common_flags.define_common_flags()
 
+# Parsing sys.argv so we can use flags by `import deepray`
+flags.FLAGS(sys.argv, known_only=True)
+if flags.FLAGS.random_seed is not None:
+  set_random_seed(flags.FLAGS.random_seed)
+
 
 def init():
   logger.debug(f"sys.argv = {sys.argv}")  # sys.argv from Horovod
+  # Parsing `distribution_strategy` this additional flag
   flags.FLAGS(sys.argv, known_only=True)
 
   gpus = tf.config.list_physical_devices("GPU")

deepray/core/trainer.py

@@ -62,13 +62,6 @@
 logger = logging_util.get_logger()
 
 
-def set_random_seed(random_seed):
-  random.seed(random_seed)  # set random seed for python
-  np.random.seed(random_seed)  # set random seed for numpy
-  tf.random.set_seed(random_seed)  # set random seed for tensorflow-cpu
-  os.environ["TF_DETERMINISTIC_OPS"] = "1"  # set random seed for tensorflow-gpu
-
-
 @keras_export("keras.Model", "keras.models.Model")
 class Trainer:
   """A model grouping layers into an object with training/inference features.
@@ -349,8 +342,6 @@ def __init__(
       logger.info("flags.FLAGS:")
       for key, value in sorted(flags.FLAGS.flag_values_dict().items()):
         logger.info(f"\t{key:25}= {value}")
-    if flags.FLAGS.random_seed is not None:
-      set_random_seed(flags.FLAGS.random_seed)
 
   def _create_counter_variable(self, init_value):
     """Helper function for counter variable creation.

deepray/utils/keras_utils.py

@@ -15,7 +15,9 @@
 
 import multiprocessing
 import os
+import random
 
+import numpy as np
 import tensorflow as tf
 from absl import logging
 from tensorflow.python import tf2
@@ -153,3 +155,11 @@ def count_params(model):
   main_print(f"{'':20} ({model_size:,})")
 
   return model_size
+
+
+def set_random_seed(random_seed):
+  random.seed(random_seed)  # set random seed for python
+  np.random.seed(random_seed)  # set random seed for numpy
+  tf.random.set_seed(random_seed)  # set random seed for tensorflow-cpu
+  # The following operation will significantly reduce the training speed, so we have disabled it.
+  # os.environ["TF_DETERMINISTIC_OPS"] = "1"  # set random seed for tensorflow-gpu

deepray/version.py

@@ -21,7 +21,7 @@
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = "0"
 _MINOR_VERSION = "21"
-_PATCH_VERSION = "96"
+_PATCH_VERSION = "97"
 
 # When building releases, we can update this value on the release branch to
 # reflect the current release candidate ('rc0', 'rc1') or, finally, the official

modelzoo/LanguageModeling/Word2Vec/dataset.py

@@ -0,0 +1,45 @@
+import tensorflow as tf
+import tqdm
+
+
+# Generates skip-gram pairs with negative sampling for a list of sequences
+# (int-encoded sentences) based on window size, number of negative samples
+# and vocabulary size.
+def generate_training_data(sequences, window_size, num_ns, vocab_size, seed):
+  # Elements of each training example are appended to these lists.
+  targets, contexts, labels = [], [], []
+
+  # Build the sampling table for `vocab_size` tokens.
+  sampling_table = tf.keras.preprocessing.sequence.make_sampling_table(vocab_size)
+
+  # Iterate over all sequences (sentences) in the dataset.
+  for sequence in tqdm.tqdm(sequences):
+    # Generate positive skip-gram pairs for a sequence (sentence).
+    positive_skip_grams, _ = tf.keras.preprocessing.sequence.skipgrams(
+      sequence, vocabulary_size=vocab_size, sampling_table=sampling_table, window_size=window_size, negative_samples=0
+    )
+
+    # Iterate over each positive skip-gram pair to produce training examples
+    # with a positive context word and negative samples.
+    for target_word, context_word in positive_skip_grams:
+      context_class = tf.expand_dims(tf.constant([context_word], dtype="int64"), 1)
+      negative_sampling_candidates, _, _ = tf.random.log_uniform_candidate_sampler(
+        true_classes=context_class,
+        num_true=1,
+        num_sampled=num_ns,
+        unique=True,
+        range_max=vocab_size,
+        seed=seed,
+        name="negative_sampling",
+      )
+
+      # Build context and label vectors (for one target word)
+      context = tf.concat([tf.squeeze(context_class, 1), negative_sampling_candidates], 0)
+      label = tf.constant([1] + [0] * num_ns, dtype="int64")
+
+      # Append each element from the training example to global lists.
+      targets.append(target_word)
+      contexts.append(context)
+      labels.append(label)
+
+  return targets, contexts, labels

modelzoo/LanguageModeling/Word2Vec/model.py

@@ -0,0 +1,81 @@
+import os
+
+import tensorflow as tf
+from absl import flags
+from tensorflow.keras import layers
+from packaging.version import parse
+from deepray.layers.dynamic_embedding import DistributedDynamicEmbedding
+from deepray.layers.embedding_variable import EmbeddingVariable
+
+if parse(tf.__version__.replace("-tf", "+tf")) < parse("2.11"):
+  from tensorflow.keras import layers as keras_layers
+  from tensorflow.keras import initializers as keras_initializers
+  from tensorflow.keras import models as keras_models
+else:
+  from tf_keras import layers as keras_layers
+  from tf_keras import initializers as keras_initializers
+  from tf_keras import models as keras_models
+
+
+class Word2Vec(tf.keras.Model):
+  def __init__(self, vocab_size, embedding_dim):
+    super(Word2Vec, self).__init__()
+    if os.getenv("USE_TF") == "1":
+      self.target_embedding = layers.Embedding(
+        vocab_size,
+        embedding_dim,
+        name="w2v_embedding",
+        embeddings_initializer=keras_initializers.Constant(value=0.1),
+      )
+      self.context_embedding = layers.Embedding(
+        vocab_size,
+        embedding_dim,
+        embeddings_initializer=keras_initializers.Constant(value=0.1),
+      )
+    elif flags.FLAGS.use_dynamic_embedding:
+      self.target_embedding = DistributedDynamicEmbedding(
+        embedding_dim=embedding_dim,
+        name="w2v_embedding",
+        key_dtype=tf.int64,
+        value_dtype=tf.float32,
+        initializer=keras_initializers.Constant(value=0.1),
+      )
+      self.context_embedding = DistributedDynamicEmbedding(
+        embedding_dim=embedding_dim,
+        key_dtype=tf.int64,
+        value_dtype=tf.float32,
+        initializer=keras_initializers.Constant(value=0.1),
+      )
+    else:
+      self.target_embedding = EmbeddingVariable(
+        embedding_dim=embedding_dim,
+        name="w2v_embedding",
+        key_dtype=tf.int64,
+        value_dtype=tf.float32,
+        storage_type="DRAM",
+        optimizer_type="Adam",
+        initializer=keras_initializers.Constant(value=0.1),
+      )
+      self.context_embedding = EmbeddingVariable(
+        embedding_dim=embedding_dim,
+        key_dtype=tf.int64,
+        value_dtype=tf.float32,
+        storage_type="DRAM",
+        optimizer_type="Adam",
+        initializer=keras_initializers.Constant(value=0.1),
+      )
+
+  def call(self, pair):
+    target, context = pair
+    # target: (batch, dummy?)  # The dummy axis doesn't exist in TF2.7+
+    # context: (batch, context)
+    if len(target.shape) == 2:
+      target = tf.squeeze(target, axis=1)
+    # target: (batch,)
+    word_emb = self.target_embedding(target)
+    # word_emb: (batch, embed)
+    context_emb = self.context_embedding(context)
+    # context_emb: (batch, context, embed)
+    dots = tf.einsum("be,bce->bc", word_emb, context_emb)
+    # dots: (batch, context)
+    return dots

modelzoo/LanguageModeling/Word2Vec/run.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -eu
+set -o pipefail
+
+batch_size=${1:-"1024"}
+learning_rate=${2:-"0.001"}
+epochs=${3:-"20"}
+
+printf -v TAG "dp_training_word2vec_gbs%d" $batch_size
+DATESTAMP=$(date +'%y%m%d%H%M%S')
+
+#Edit to save logs & checkpoints in a different directory
+RESULTS_DIR=${PWD}
+LOGFILE=$RESULTS_DIR/$TAG.$DATESTAMP.log
+mkdir -m 777 -p $RESULTS_DIR
+printf "Saving checkpoints to %s\n" "$RESULTS_DIR"
+printf "Logs written to %s\n" "$LOGFILE"
+
+set -x
+export CUDA_VISIBLE_DEVICES=0
+export USE_TF=1
+export USE_DP=1
+python train.py \
+    --use_dynamic_embedding=True \
+    --use_custom_training_loop=False \
+    --random_seed=42 \
+    --batch_size=$batch_size \
+    --steps_per_execution=1 \
+    --run_eagerly=False \
+    --learning_rate=$learning_rate \
+    --epochs=$epochs \
+    --model_dir=${RESULTS_DIR} \
+    $@
+# |& tee $LOGFILE
+set +x