Internal change

PiperOrigin-RevId: 339565966

official/vision/beta/configs/backbones_3d.py

@@ -34,8 +34,9 @@ class ResNet3DBlock(hyperparams.Config):
 class ResNet3D(hyperparams.Config):
   """ResNet config."""
   model_id: int = 50
-  stem_temporal_conv_stride: int = 2
-  stem_temporal_pool_stride: int = 2
+  stem_conv_temporal_kernel_size: int = 5
+  stem_conv_temporal_stride: int = 2
+  stem_pool_temporal_stride: int = 2
   block_specs: Tuple[ResNet3DBlock, ...] = ()
 
 

official/vision/beta/configs/experiments/video_classification/k400_3d-resnet50_tpu.yaml

+# 3D ResNet-50 video classification on Kinetics-400. 75.3% top-1 and 91.2% top-5 accuracy.
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+task:
+  model:
+    backbone:
+      resnet_3d:
+        block_specs: !!python/tuple
+        - temporal_kernel_sizes: !!python/tuple
+          - 3
+          - 3
+          - 3
+          temporal_strides: 1
+          use_self_gating: true
+        - temporal_kernel_sizes: !!python/tuple
+          - 3
+          - 1
+          - 3
+          - 1
+          temporal_strides: 1
+          use_self_gating: true
+        - temporal_kernel_sizes: !!python/tuple
+          - 3
+          - 1
+          - 3
+          - 1
+          - 3
+          - 1
+          temporal_strides: 1
+          use_self_gating: true
+        - temporal_kernel_sizes: !!python/tuple
+          - 1
+          - 3
+          - 1
+          temporal_strides: 1
+          use_self_gating: true
+        model_id: 50
+        stem_conv_temporal_kernel_size: 5
+        stem_conv_temporal_stride: 2
+        stem_pool_temporal_stride: 2
+    dropout_rate: 0.5
+  train_data:
+    global_batch_size: 1024
+    dtype: 'bfloat16'
+    shuffle_buffer_size: 1024
+  validation_data:
+    global_batch_size: 32
+    dtype: 'bfloat16'
+    drop_remainder: false
+trainer:
+  optimizer_config:
+    learning_rate:
+      cosine:
+        initial_learning_rate: 0.8
+        decay_steps: 42104
+    warmup:
+      linear:
+        warmup_steps: 1053
+  train_steps: 42104
+  steps_per_loop: 500
+  summary_interval: 500
+  validation_interval: 500

official/vision/beta/configs/experiments/video_classification/k400_slowonly8x8_tpu.yaml

+# SlowOnly video classification on Kinetics-400. Expected performance to be updated.
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+task:
+  model:
+    backbone:
+      resnet_3d:
+        block_specs: !!python/tuple
+        - temporal_kernel_sizes: !!python/tuple
+          - 1
+          - 1
+          - 1
+          temporal_strides: 1
+          use_self_gating: false
+        - temporal_kernel_sizes: !!python/tuple
+          - 1
+          - 1
+          - 1
+          - 1
+          temporal_strides: 1
+          use_self_gating: false
+        - temporal_kernel_sizes: !!python/tuple
+          - 3
+          - 3
+          - 3
+          - 3
+          - 3
+          - 3
+          temporal_strides: 1
+          use_self_gating: false
+        - temporal_kernel_sizes: !!python/tuple
+          - 3
+          - 3
+          - 3
+          temporal_strides: 1
+          use_self_gating: false
+        model_id: 50
+        stem_conv_temporal_kernel_size: 1
+        stem_conv_temporal_stride: 1
+        stem_pool_temporal_stride: 1
+    dropout_rate: 0.5
+  train_data:
+    feature_shape: !!python/tuple
+    - 8
+    - 224
+    - 224
+    - 3
+    temporal_stride: 8
+    global_batch_size: 1024
+    dtype: 'bfloat16'
+    shuffle_buffer_size: 1024
+  validation_data:
+    global_batch_size: 32
+    dtype: 'bfloat16'
+    drop_remainder: false
+trainer:
+  optimizer_config:
+    learning_rate:
+      cosine:
+        initial_learning_rate: 0.8
+        decay_steps: 42104
+    warmup:
+      linear:
+        warmup_steps: 1053
+  train_steps: 42104
+  steps_per_loop: 500
+  summary_interval: 500
+  validation_interval: 500

official/vision/beta/modeling/backbones/resnet_3d.py

@@ -49,6 +49,7 @@ class ResNet3D(tf.keras.Model):
                temporal_kernel_sizes: List[Tuple[int]],
                use_self_gating: List[int] = None,
                input_specs=layers.InputSpec(shape=[None, None, None, None, 3]),
+               stem_conv_temporal_kernel_size=5,
                stem_conv_temporal_stride=2,
                stem_pool_temporal_stride=2,
                activation='relu',
@@ -70,6 +71,8 @@ class ResNet3D(tf.keras.Model):
       use_self_gating: a list of booleans to specify applying self-gating module
         or not in each block group. If None, self-gating is not applied.
       input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
+      stem_conv_temporal_kernel_size: `int` temporal kernel size for the first
+        conv layer.
       stem_conv_temporal_stride: `int` temporal stride for the first conv layer.
       stem_pool_temporal_stride: `int` temporal stride for the first pool layer.
       activation: `str` name of the activation function.
@@ -88,6 +91,7 @@ class ResNet3D(tf.keras.Model):
     self._temporal_strides = temporal_strides
     self._temporal_kernel_sizes = temporal_kernel_sizes
     self._input_specs = input_specs
+    self._stem_conv_temporal_kernel_size = stem_conv_temporal_kernel_size
     self._stem_conv_temporal_stride = stem_conv_temporal_stride
     self._stem_pool_temporal_stride = stem_pool_temporal_stride
     self._use_self_gating = use_self_gating
@@ -113,7 +117,7 @@ class ResNet3D(tf.keras.Model):
     # Build stem.
     x = layers.Conv3D(
         filters=64,
-        kernel_size=[5, 7, 7],
+        kernel_size=[stem_conv_temporal_kernel_size, 7, 7],
         strides=[stem_conv_temporal_stride, 2, 2],
         use_bias=False,
         padding='same',
@@ -239,6 +243,7 @@ class ResNet3D(tf.keras.Model):
         'model_id': self._model_id,
         'temporal_strides': self._temporal_strides,
         'temporal_kernel_sizes': self._temporal_kernel_sizes,
+        'stem_conv_temporal_kernel_size': self._stem_conv_temporal_kernel_size,
         'stem_conv_temporal_stride': self._stem_conv_temporal_stride,
         'stem_pool_temporal_stride': self._stem_pool_temporal_stride,
         'use_self_gating': self._use_self_gating,

official/vision/beta/modeling/backbones/resnet_3d_test.py

@@ -65,6 +65,7 @@ class ResNet3DTest(parameterized.TestCase, tf.test.TestCase):
         temporal_strides=[1, 1, 1, 1],
         temporal_kernel_sizes=[(3, 3, 3), (3, 1, 3, 1), (3, 1, 3, 1, 3, 1),
                                (1, 3, 1)],
+        stem_conv_temporal_kernel_size=5,
         stem_conv_temporal_stride=2,
         stem_pool_temporal_stride=2,
         use_self_gating=None,