Prepare for upcoming keras initializer change.

PiperOrigin-RevId: 451476877

official/projects/s3d/modeling/inception_utils.py

@@ -17,6 +17,7 @@ from typing import Callable, Dict, Optional, Sequence, Set, Text, Tuple, Type, U
 
 import tensorflow as tf
 
+from official.modeling import tf_utils
 from official.projects.s3d.modeling import net_utils
 from official.vision.modeling.layers import nn_blocks_3d
 
@@ -126,7 +127,7 @@ def inception_v1_stem_cells(
       strides=[2, 2, 2],
       padding='same',
       use_bias=False,
-      kernel_initializer=kernel_initializer,
+      kernel_initializer=tf_utils.clone_initializer(kernel_initializer),
       kernel_regularizer=kernel_regularizer,
       name=layer_naming_fn(end_point))(
           inputs)
@@ -161,7 +162,7 @@ def inception_v1_stem_cells(
       kernel_size=[1, 1, 1],
       padding='same',
       use_bias=False,
-      kernel_initializer=kernel_initializer,
+      kernel_initializer=tf_utils.clone_initializer(kernel_initializer),
       kernel_regularizer=kernel_regularizer,
       name=layer_naming_fn(end_point))(
           net)
@@ -191,7 +192,7 @@ def inception_v1_stem_cells(
       norm_momentum=norm_momentum,
       norm_epsilon=norm_epsilon,
       temporal_conv_initializer=temporal_conv_initializer,
-      kernel_initializer=kernel_initializer,
+      kernel_initializer=tf_utils.clone_initializer(kernel_initializer),
       kernel_regularizer=kernel_regularizer,
       name=layer_naming_fn(end_point))(
           net)
@@ -330,7 +331,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             kernel_size=[1, 1, 1],
             padding='same',
             use_bias=False,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer),
         # norm
         dict(
@@ -349,7 +351,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             kernel_size=[1, 1, 1],
             padding='same',
             use_bias=False,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer),
         # norm
         dict(
@@ -371,7 +374,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             norm_momentum=self._norm_momentum,
             norm_epsilon=self._norm_epsilon,
             temporal_conv_initializer=self._temporal_conv_initializer,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer),
     ]
     branch_2_params = [
@@ -381,7 +385,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             kernel_size=[1, 1, 1],
             padding='same',
             use_bias=False,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer),
         # norm
         dict(
@@ -403,7 +408,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             norm_momentum=self._norm_momentum,
             norm_epsilon=self._norm_epsilon,
             temporal_conv_initializer=self._temporal_conv_initializer,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer)
     ]
     branch_3_params = [
@@ -413,7 +419,8 @@ class InceptionV1CellLayer(tf.keras.layers.Layer):
             kernel_size=[1, 1, 1],
             padding='same',
             use_bias=False,
-            kernel_initializer=self._kernel_initializer,
+            kernel_initializer=tf_utils.clone_initializer(
+                self._kernel_initializer),
             kernel_regularizer=self._kernel_regularizer),
         # norm
         dict(

official/projects/s3d/modeling/net_utils.py

@@ -16,6 +16,7 @@
 from typing import Any, Text, Sequence, Union
 
 import tensorflow as tf
+from official.modeling import tf_utils
 
 WEIGHT_INITIALIZER = {
     'Xavier': tf.keras.initializers.GlorotUniform,
@@ -94,7 +95,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[self._kernel_size] * 3,
               strides=self._strides,
               dilation_rate=self._rates,
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
     elif self._conv_type == '2d':
       conv_layer_params.append(
@@ -103,7 +105,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[1, self._kernel_size, self._kernel_size],
               strides=[1, self._strides[1], self._strides[2]],
               dilation_rate=[1, self._rates[1], self._rates[2]],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
     elif self._conv_type == '1+2d':
       channels_in = input_shape[self._channel_axis]
@@ -113,7 +116,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[self._kernel_size, 1, 1],
               strides=[self._strides[0], 1, 1],
               dilation_rate=[self._rates[0], 1, 1],
-              kernel_initializer=self._temporal_conv_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._temporal_conv_initializer),
           ))
       conv_layer_params.append(
           dict(
@@ -121,7 +125,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[1, self._kernel_size, self._kernel_size],
               strides=[1, self._strides[1], self._strides[2]],
               dilation_rate=[1, self._rates[1], self._rates[2]],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
     elif self._conv_type == '2+1d':
       conv_layer_params.append(
@@ -130,7 +135,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[1, self._kernel_size, self._kernel_size],
               strides=[1, self._strides[1], self._strides[2]],
               dilation_rate=[1, self._rates[1], self._rates[2]],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
       conv_layer_params.append(
           dict(
@@ -138,7 +144,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[self._kernel_size, 1, 1],
               strides=[self._strides[0], 1, 1],
               dilation_rate=[self._rates[0], 1, 1],
-              kernel_initializer=self._temporal_conv_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._temporal_conv_initializer),
           ))
     elif self._conv_type == '1+1+1d':
       conv_layer_params.append(
@@ -147,7 +154,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[1, 1, self._kernel_size],
               strides=[1, 1, self._strides[2]],
               dilation_rate=[1, 1, self._rates[2]],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
       conv_layer_params.append(
           dict(
@@ -155,7 +163,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[1, self._kernel_size, 1],
               strides=[1, self._strides[1], 1],
               dilation_rate=[1, self._rates[1], 1],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
       conv_layer_params.append(
           dict(
@@ -163,7 +172,8 @@ class ParameterizedConvLayer(tf.keras.layers.Layer):
               kernel_size=[self._kernel_size, 1, 1],
               strides=[self._strides[0], 1, 1],
               dilation_rate=[self._rates[0], 1, 1],
-              kernel_initializer=self._kernel_initializer,
+              kernel_initializer=tf_utils.clone_initializer(
+                  self._kernel_initializer),
           ))
     else:
       raise ValueError('Unsupported conv_type: {}'.format(self._conv_type))