Add Imagenette dataset to dataset_lib

init2winit/dataset_lib/datasets.py

@@ -23,6 +23,7 @@
 from init2winit.dataset_lib import fake_dataset
 from init2winit.dataset_lib import fastmri_dataset
 from init2winit.dataset_lib import imagenet_dataset
+from init2winit.dataset_lib import imagenette_dataset
 from init2winit.dataset_lib import librispeech
 from init2winit.dataset_lib import lm1b_v2
 # We get TF v2 eager execution error if we import fineweb_edu_10b
@@ -112,6 +113,12 @@
         imagenet_dataset.METADATA,
         imagenet_dataset.get_fake_batch,
     ),
+    'imagenette': _Dataset(
+        imagenette_dataset.get_imagenette,
+        imagenette_dataset.DEFAULT_HPARAMS,
+        imagenette_dataset.METADATA,
+        imagenette_dataset.get_fake_batch,
+    ),
     'translate_wmt': _Dataset(
         translate_wmt.get_translate_wmt,
         translate_wmt.DEFAULT_HPARAMS,

init2winit/dataset_lib/imagenet_dataset.py

@@ -67,6 +67,7 @@ class PreprocessForTrainTransform(grain.RandomMapTransform):
   use_randaug: bool
   randaug_magnitude: int
   randaug_num_layers: int
+  num_classes: int = imagenet_preprocessing.NUM_CLASSES
 
   def random_map(self, features, seed):
     inputs = imagenet_preprocessing.preprocess_for_train(
@@ -80,7 +81,7 @@ def random_map(self, features, seed):
         randaug_magnitude=self.randaug_magnitude,
         randaug_num_layers=self.randaug_num_layers,
     )
-    targets = tf.one_hot(features['label'], imagenet_preprocessing.NUM_CLASSES)
+    targets = tf.one_hot(features['label'], self.num_classes)
     result = {'inputs': inputs, 'targets': targets, 'weights': 1}
     for k in grain.META_FEATURES:
       if k in features:
@@ -95,12 +96,13 @@ class PreprocessForEvalTransform(grain.MapTransform):
   dtype: Any
   image_size: int
   include_example_keys: bool
+  num_classes: int = imagenet_preprocessing.NUM_CLASSES
 
   def map(self, features):
     inputs = imagenet_preprocessing.preprocess_for_eval(
         features['image'], self.dtype, self.image_size
     )
-    targets = tf.one_hot(features['label'], imagenet_preprocessing.NUM_CLASSES)
+    targets = tf.one_hot(features['label'], self.num_classes)
     result = {
         'inputs': inputs,
         'targets': targets,
@@ -155,6 +157,7 @@ def load_split_grain(
     dtype=tf.float32,
     shuffle_rng=None,
     tfds_dataset_name='imagenet2012:5.*.*',
+    num_classes=imagenet_preprocessing.NUM_CLASSES,
 ):
   """Uses Grain to load the data. See documentation for `load_split`."""
   # Grain starts counting at 1
@@ -181,6 +184,7 @@ def load_split_grain(
             use_randaug=hps.use_randaug,
             randaug_magnitude=hps.randaug.magnitude,
             randaug_num_layers=hps.randaug.num_layers,
+            num_classes=num_classes,
         )
     ]
   elif split == 'eval_train':
@@ -192,6 +196,7 @@ def load_split_grain(
             dtype=dtype,
             image_size=image_size,
             include_example_keys=hps.get('include_example_keys'),
+            num_classes=num_classes,
         ),
     ]
   else:
@@ -201,6 +206,7 @@ def load_split_grain(
             dtype=dtype,
             image_size=image_size,
             include_example_keys=hps.get('include_example_keys'),
+            num_classes=num_classes,
         ),
         transforms.CacheTransform(),
     ]
@@ -257,7 +263,9 @@ def load_split(
     dtype=tf.float32,
     image_size=224,
     shuffle_rng=None,
-    tfds_dataset_name='imagenet2012:5.*.*'):  # pyformat: disable
+    tfds_dataset_name='imagenet2012:5.*.*',
+    num_classes=imagenet_preprocessing.NUM_CLASSES,
+):  # pyformat: disable
   """Creates a split from the ImageNet dataset using TensorFlow Datasets.
 
   The dataset returned by this function will repeat forever if split == 'train',
@@ -281,6 +289,7 @@ def load_split(
       'train'`.
     tfds_dataset_name: The name of the dataset to load from TFDS. Used to reuse
       this same logic for imagenet-v2.
+    num_classes: The number of classes to one-hot encode target labels.
 
   Returns:
     A `tf.data.Dataset`.
@@ -326,12 +335,16 @@ def decode_example(example_index, example):
           hps.use_randaug,
           hps.randaug.magnitude,
           hps.randaug.num_layers,
+          num_classes=num_classes,
       )
       example_dict = transform.random_map(example, preprocess_rng)
 
     else:
       transform = PreprocessForEvalTransform(
-          dtype, image_size, hps.get('include_example_keys')
+          dtype,
+          image_size,
+          hps.get('include_example_keys'),
+          num_classes=num_classes,
       )
       example_dict = transform.map(example)
 
@@ -393,7 +406,15 @@ def mixup_batch(batch_index, batch):
   return ds
 
 
-def get_imagenet(shuffle_rng, batch_size, eval_batch_size, hps, global_step=0):
+def get_imagenet(
+    shuffle_rng,
+    batch_size,
+    eval_batch_size,
+    hps,
+    global_step=0,
+    tfds_dataset_name='imagenet2012:5.*.*',
+    num_classes=imagenet_preprocessing.NUM_CLASSES,
+):
   """Data generators for imagenet."""
   per_host_batch_size = batch_size // jax.process_count()
   per_host_eval_batch_size = eval_batch_size // jax.process_count()
@@ -414,6 +435,8 @@ def get_imagenet(shuffle_rng, batch_size, eval_batch_size, hps, global_step=0):
       image_size=image_size,
       shuffle_rng=shuffle_rng,
       global_step=global_step,
+      tfds_dataset_name=tfds_dataset_name,
+      num_classes=num_classes,
   )
   train_ds = tfds.as_numpy(train_ds)
   logging.info('Loading eval_train split')
@@ -423,6 +446,8 @@ def get_imagenet(shuffle_rng, batch_size, eval_batch_size, hps, global_step=0):
       hps=hps,
       image_size=image_size,
       global_step=global_step,
+      tfds_dataset_name=tfds_dataset_name,
+      num_classes=num_classes,
   )
   eval_train_ds = tfds.as_numpy(eval_train_ds)
   logging.info('Loading eval split')
@@ -432,18 +457,21 @@ def get_imagenet(shuffle_rng, batch_size, eval_batch_size, hps, global_step=0):
       hps=hps,
       image_size=image_size,
       global_step=global_step,
+      tfds_dataset_name=tfds_dataset_name,
+      num_classes=num_classes,
   )
   validation_ds = tfds.as_numpy(validation_ds)
 
   test_ds = None
-  if hps.use_imagenetv2_test:
+  if hps.get('use_imagenetv2_test'):
     test_ds = load_split_fn(
         per_host_eval_batch_size,
         'test',
         hps=hps,
         image_size=image_size,
         tfds_dataset_name='imagenet_v2/matched-frequency',
         global_step=global_step,
+        num_classes=num_classes,
     )
     test_ds = tfds.as_numpy(test_ds)
 

init2winit/dataset_lib/imagenette_dataset.py

@@ -0,0 +1,71 @@
+# coding=utf-8
+# Copyright 2026 The init2winit Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Imagenette input pipeline with ImageNet-style preprocessing.
+
+Imagenette is a subset of 10 easily classified classes from ImageNet,
+created by Jeremy Howard / fastai. It uses the same preprocessing as ImageNet
+(JPEG decode, random crop, resize to 224x224, ImageNet mean/std normalization).
+
+TFDS dataset: imagenette/full-size-v2
+  - Train: 9,469 images
+  - Validation: 3,925 images
+  - No test split
+"""
+
+from init2winit.dataset_lib import imagenet_dataset
+from ml_collections.config_dict import config_dict
+
+NUM_CLASSES = 10
+
+DEFAULT_HPARAMS = config_dict.ConfigDict(
+    dict(
+        input_shape=(224, 224, 3),
+        output_shape=(NUM_CLASSES,),
+        train_size=9469,
+        valid_size=3925,
+        test_size=0,
+        crop='random',  # options are: {"random", "inception", "center"}
+        random_flip=True,
+        use_mixup=False,
+        mixup={'alpha': 0.5},
+        use_randaug=False,
+        randaug={'magnitude': 15, 'num_layers': 2},
+        use_grain=False,
+    )
+)
+
+METADATA = {
+    'apply_one_hot_in_loss': False,
+}
+
+TFDS_DATASET_NAME = 'imagenette/full-size-v2:1.*.*'
+
+get_fake_batch = imagenet_dataset.get_fake_batch
+
+
+def get_imagenette(
+    shuffle_rng, batch_size, eval_batch_size, hps, global_step=0
+):
+  """Data generators for Imagenette."""
+  return imagenet_dataset.get_imagenet(
+      shuffle_rng=shuffle_rng,
+      batch_size=batch_size,
+      eval_batch_size=eval_batch_size,
+      hps=hps,
+      global_step=global_step,
+      tfds_dataset_name=TFDS_DATASET_NAME,
+      num_classes=NUM_CLASSES,
+  )