Fix imagenet config.

PiperOrigin-RevId: 494734670 Change-Id: I68a2ddefabd7f6d9fdfc654bd17195f059e5e954
google-deepmind · Dec 13, 2022 · d109541 · d109541
1 parent a5298e6
commit d109541
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 24 deletions.
diff --git a/experiments_jax/configs/pretrain_imagenet.py b/experiments_jax/configs/pretrain_imagenet.py
@@ -51,18 +51,19 @@ def get_config() -> ml_collections.ConfigDict:
       'experiment': {
           'resume_from_checkpoint_path': None,
           'stream': {
-              'ctor':
-                  streams.FilteredStream,
-              'stream_ctor':
-                  nevis_stream.NevisStream,
-              'supported_task_kinds': [tasks.TaskKind.CLASSIFICATION],
-              'stream_variant':
-                  nevis_stream.NevisStreamVariant.IMAGENET_ONLY,
-              'predict_event_splits': (
-                  nevis_stream.Split.DEV,
-                  nevis_stream.Split.DEV_TEST,
-                  nevis_stream.Split.TEST,
-              ),
+              'ctor': streams.FilteredStream,
+              'kwargs': {
+                  'stream_ctor':
+                      nevis_stream.NevisStream,
+                  'supported_task_kinds': [tasks.TaskKind.CLASSIFICATION],
+                  'stream_variant':
+                      nevis_stream.NevisStreamVariant.IMAGENET_ONLY,
+                  'predict_event_splits': (
+                      nevis_stream.Split.DEV,
+                      nevis_stream.Split.DEV_TEST,
+                      nevis_stream.Split.TEST,
+                  ),
+              },
           },
           'learner': {
               'learner_builder': finetuning_learner.learner_builder,

diff --git a/experiments_torch/configs/pretrain_imagenet.py b/experiments_torch/configs/pretrain_imagenet.py
@@ -48,18 +48,19 @@ def get_config() -> ml_collections.ConfigDict:
       'experiment': {
           'resume_from_checkpoint_path': None,
           'stream': {
-              'ctor':
-                  streams.FilteredStream,
-              'stream_ctor':
-                  nevis_stream.NevisStream,
-              'supported_task_kinds': [tasks.TaskKind.CLASSIFICATION],
-              'stream_variant':
-                  nevis_stream.NevisStreamVariant.IMAGENET_ONLY,
-              'predict_event_splits': (
-                  nevis_stream.Split.DEV,
-                  nevis_stream.Split.DEV_TEST,
-                  nevis_stream.Split.TEST,
-              ),
+              'ctor': streams.FilteredStream,
+              'kwargs': {
+                  'stream_ctor':
+                      nevis_stream.NevisStream,
+                  'supported_task_kinds': [tasks.TaskKind.CLASSIFICATION],
+                  'stream_variant':
+                      nevis_stream.NevisStreamVariant.IMAGENET_ONLY,
+                  'predict_event_splits': (
+                      nevis_stream.Split.DEV,
+                      nevis_stream.Split.DEV_TEST,
+                      nevis_stream.Split.TEST,
+                  ),
+              },
           },
           'learner': {
               'learner_builder': finetuning_learner.learner_builder,