Internal change

PiperOrigin-RevId: 397350737

official/vision/beta/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml

-# ResNet-50 ImageNet classification. 78.1% top-1 and 93.9% top-5 accuracy.
 runtime:
   distribution_strategy: 'tpu'
   mixed_precision_dtype: 'bfloat16'
@@ -10,8 +9,6 @@ task:
       type: 'resnet'
       resnet:
         model_id: 50
-    norm_activation:
-      activation: 'swish'
   losses:
     l2_weight_decay: 0.0001
     one_hot: true
@@ -28,7 +25,7 @@ task:
     dtype: 'bfloat16'
     drop_remainder: false
 trainer:
-  train_steps: 62400
+  train_steps: 28080
   validation_steps: 13
   validation_interval: 312
   steps_per_loop: 312
@@ -40,10 +37,10 @@ trainer:
       sgd:
         momentum: 0.9
     learning_rate:
-      type: 'cosine'
-      cosine:
-        initial_learning_rate: 1.6
-        decay_steps: 62400
+      type: 'stepwise'
+      stepwise:
+        boundaries: [9360, 18720, 24960]
+        values: [1.6, 0.16, 0.016, 0.0016]
     warmup:
       type: 'linear'
       linear: