Merge pull request #2812 from nouiz:imagenet_benchmark

PiperOrigin-RevId: 510490709
google · Feb 17, 2023 · 018725f · 018725f
2 parents 5d4040a + da77531
commit 018725f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/examples/imagenet/imagenet_fake_data_benchmark.py b/examples/imagenet/imagenet_fake_data_benchmark.py
@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Benchmark for the ImageNet example using fake data for quick perf results."""
+"""Benchmark for the ImageNet example using fake data for quick perf results.
+
+This script doesn't need the dataset, but it needs the dataset metadata.
+That can be fetched with the script `flax/tests/download_dataset_metadata.sh`.
+"""
 
 import pathlib
 import time

diff --git a/examples/imagenet/train.py b/examples/imagenet/train.py
@@ -289,7 +289,7 @@ def train_and_evaluate(config: ml_collections.ConfigDict,
       dataset_builder.info.splits['train'].num_examples // config.batch_size
   )
 
-  if config.num_train_steps == -1:
+  if config.num_train_steps <= 0:
     num_steps = int(steps_per_epoch * config.num_epochs)
   else:
     num_steps = config.num_train_steps