clean up ogbg

Author: priyakasimbeg

algoperf/workloads/ogbg/input_pipeline.py

@@ -148,22 +148,10 @@ def _get_batch_iterator(dataset_iter, global_batch_size, num_shards=None):
     weights_shards.append(weights)
 
     if count == num_shards:
-      # yield {
-      #     'inputs': jraph.batch(graphs_shards),
-      #     'targets': np.vstack(labels_shards),
-      #     'weights': np.vstack(weights_shards)
-      # }
-
-      def f(x):
-        return jax.tree.map(lambda *vals: np.concatenate(vals, axis=0), x[0], *x[1:])
-      
-      graphs_shards = f(graphs_shards)
-      labels_shards = f(labels_shards)
-      weights_shards = f(weights_shards)
       yield {
-          'inputs': graphs_shards,
-          'targets': labels_shards,
-          'weights': weights_shards,
+          'inputs': jraph.batch(graphs_shards),
+          'targets': np.vstack(labels_shards),
+          'weights': np.vstack(weights_shards)
       }
 
       count = 0

algoperf/workloads/ogbg/ogbg_jax/models.py

@@ -79,8 +79,7 @@ def __call__(self, graph, train):
               self.hidden_dims, dropout=dropout, activation_fn=activation_fn),
           update_global_fn=_make_mlp(
               self.hidden_dims, dropout=dropout, activation_fn=activation_fn))
-      # jax.debug.print(str(graph))
-    
+
       graph = net(graph)
 
     # Map globals to represent the final result

algoperf/workloads/ogbg/ogbg_jax/workload.py

@@ -108,11 +108,7 @@ def _eval_metric(self, labels, logits, masks):
     return metrics.EvalMetrics.single_from_model_output(
         loss=loss['per_example'], logits=logits, labels=labels, mask=masks)
 
-  # @functools.partial(
-  #     jax.pmap,
-  #     axis_name='batch',
-  #     in_axes=(None, 0, 0, 0, None),
-  #     static_broadcasted_argnums=(0,))
+
   @functools.partial(
     jax.jit,
     in_shardings=(sharding_utils.get_replicated_sharding(),
@@ -130,8 +126,6 @@ def _normalize_eval_metrics(
                                                    Any]) -> Dict[str, float]:
     """Normalize eval metrics."""
     del num_examples
-    # total_metrics = total_metrics.reduce()
-    print(total_metrics)
     return {k: float(v) for k, v in total_metrics.compute().items()}
 
 

docker/build_docker_images.sh

@@ -14,7 +14,7 @@ do
 done
 
 # Artifact repostiory
-ARTIFACT_REPO="europe-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo"
+ARTIFACT_REPO="europe-west-4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo"
 
 if [[ -z ${GIT_BRANCH+x} ]]
 then 

docker/scripts/startup.sh

@@ -293,7 +293,6 @@ if [[ ! -z ${SUBMISSION_PATH+x} ]]; then
         --workload=${WORKLOAD} \
         --submission_path=${SUBMISSION_PATH}  \
         --data_dir=${DATA_DIR} \
-        --num_tuning_trials=1  \
         --experiment_dir=${EXPERIMENT_DIR}  \
         --experiment_name=${EXPERIMENT_NAME} \
         --overwrite=${OVERWRITE} \

prize_qualification_baselines/self_tuning/pytorch_nadamw_full_budget.py

@@ -1,5 +1,6 @@
 """Submission file for an NAdamW optimizer with warmup+cosine LR in PyTorch."""
 
+import collections
 import math
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
@@ -24,6 +25,7 @@
     "weight_decay": 0.08121616522670176,
     "warmup_factor": 0.02
 }
+HPARAMS = collections.namedtuple('Hyperparameters', HPARAMS.keys())(**HPARAMS)
 
 
 # Modified from github.com/pytorch/pytorch/blob/v1.12.1/torch/optim/adamw.py.

prize_qualification_baselines/self_tuning/pytorch_nadamw_target_setting.py

@@ -1,5 +1,6 @@
 """Submission file for an NAdamW optimizer with warmup+cosine LR in PyTorch."""
 
+import collections
 import math
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
@@ -24,6 +25,7 @@
     "weight_decay": 0.08121616522670176,
     "warmup_factor": 0.02
 }
+HPARAMS = collections.namedtuple('Hyperparameters', HPARAMS.keys())(**HPARAMS)
 
 
 # Modified from github.com/pytorch/pytorch/blob/v1.12.1/torch/optim/adamw.py.