Change the tf-models-official to avoid a class conflict.

PiperOrigin-RevId: 698999199

Author: pichuan

deepvariant/dv_config.py

@@ -360,8 +360,9 @@ def get_config(config_name: str) -> ml_collections.ConfigDict:
   # Stop training when this many consecutive evaluations yield no improvement.
   config.early_stopping_patience = 250
 
-  # Weight decay of optimizer
+  # Optimizer params
   config.optimizer_weight_decay = 0.0
+  config.ema_momentum = 0.99
 
   # An 'iter' refers to a group of train/tune steps run in succession.
   config.steps_per_iter = 128

deepvariant/train.py

@@ -49,7 +49,7 @@
 from deepvariant import keras_modeling
 from official.modeling import optimization
 
-_CHECKPOINT_OPTIONS = tf.train.CheckpointOptions(enable_async=True)
+_CHECKPOINT_OPTIONS = tf.train.CheckpointOptions()
 
 _LEADER = flags.DEFINE_string(
     'leader',

run-prereq.sh

@@ -289,6 +289,15 @@ sudo -H NEEDRESTART_MODE=a apt-get install "${APT_ARGS[@]}" libssl-dev libcurl4-
 # for the debruijn graph
 sudo -H NEEDRESTART_MODE=a apt-get install "${APT_ARGS[@]}" libboost-graph-dev > /dev/null
 
+# Pin tf-models-official back to 2.11.6 to be closer to
+# ${DV_GCP_OPTIMIZED_TF_WHL_VERSION} (which is 2.11.0).
+# This is to avoid the issue:
+# ValueError: Addons>LAMB has already been registered to <class 'tensorflow_addons.optimizers.lamb.LAMB'>
+# However, it's important that the protobuf pinning happens after this!
+# TODO: Remove this later once the first dependency can be changed
+#                to ${DV_GCP_OPTIMIZED_TF_WHL_VERSION}.
+pip3 install "${PIP_ARGS[@]}"  "tf-models-official==2.11.6"
+
 # Just being safe, pin protobuf's version one more time.
 pip3 install "${PIP_ARGS[@]}" 'protobuf==3.13.0'