Added time log for each step and removed Bfloat16cast

oper/run_graphcast.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/usr/bin/env -S python3 -u
 
 '''
 Description: Script to call the graphcast model using gdas products
@@ -9,6 +9,7 @@
 '''
 import os
 import argparse
+from time import time
 from datetime import timedelta
 import dataclasses
 import functools
@@ -163,9 +164,13 @@ def construct_wrapped_graphcast(model_config, task_config):
             # from/to float32 to/from BFloat16.
             predictor = casting.Bfloat16Cast(predictor)
 
-            # Modify inputs/outputs to `casting.Bfloat16Cast` so the casting to/from
-            # BFloat16 happens after applying normalization to the inputs/targets.
-            predictor = normalization.InputsAndResiduals(predictor, diffs_stddev_by_level=self.diffs_stddev_by_level, mean_by_level=self.mean_by_level, stddev_by_level=self.stddev_by_level,)
+            # Applying normalization to the inputs/targets.
+            predictor = normalization.InputsAndResiduals(
+                predictor, 
+                diffs_stddev_by_level=self.diffs_stddev_by_level, 
+                mean_by_level=self.mean_by_level, 
+                stddev_by_level=self.stddev_by_level,
+            )
 
             # Wraps everything so the one-step model can produce trajectories.
             predictor = autoregressive.Predictor(predictor, gradient_checkpointing=True,)
@@ -176,8 +181,11 @@ def run_forward(model_config, task_config, inputs, targets_template, forcings,):
             predictor = construct_wrapped_graphcast(model_config, task_config)
             return predictor(inputs, targets_template=targets_template, forcings=forcings,)
 
+        t0 = time()
         jax.jit(self._with_configs(run_forward.init))
         self.model = self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply))))
+        elapsed_time = time() - t0
+        print(f"Elapsed time for compiling the model: {elapsed_time} seconds")
 
 
     def get_predictions(self):
@@ -280,11 +288,30 @@ def upload_to_s3(self, keep_data):
     args = parser.parse_args()
     runner = GraphCastModel(args.weights, args.input, args.case_name, args.config, args.output, int(args.pressure), int(args.length))
 
+    t0 = time()
     runner.load_pretrained_model()
+    elapsed_time = time() - t0
+    print(f"Elapsed time for loading model: {elapsed_time} seconds")
+
+    t0 = time()
     runner.load_gdas_data()
+    elapsed_time = time() - t0
+    print(f"Elapsed time for loading input data: {elapsed_time} seconds")
+
+    t0 = time()
     runner.extract_inputs_targets_forcings()
+    elapsed_time = time() - t0
+    print(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds")
+
+    t0 = time()
     runner.load_normalization_stats()
+    elapsed_time = time() - t0
+    print(f"Elapsed time for loading normalization stats: {elapsed_time} seconds")
+
+    t0 = time()
     runner.get_predictions()
+    elapsed_time = time() - t0
+    print(f"Elapsed time for running the model: {elapsed_time} seconds")
 
     upload_data = args.upload.lower() == "yes"
     keep_data = args.keep.lower() == "yes"