image-text proj update (#55)

Co-authored-by: Lucas Beyer <[email protected]>
Co-authored-by: Xiaohua Zhai <[email protected]>

Author: 3 people

big_vision/configs/proj/image_text/common.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# Copyright 2023 Big Vision Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -61,7 +61,7 @@ def _drop_no_imagenet(f):
         'class_names_dataset_name': 'imagenet2012',
         'pp_img': lambda sz: (
             _square875(sz) + '|pad_to_shape(inkey="real_label", outkey="label", shape=[10], pad_value=-1)|keep("label", "image")'),  # pylint: disable=line-too-long
-        'filter_fn': _drop_no_real_label,
+        'pre_filter_fn': _drop_no_real_label,
     },
     'imagenet_v2': {'class_names': 'clip'},
     'imagenet_a': {
@@ -75,16 +75,16 @@ def _drop_no_imagenet(f):
 }
 
 
-def get_disclf(sz, *, log_steps, pp_txt=None, dataset_names=('imagenet2012',)):
+def get_disclf(sz, *, pp_txt=None, dataset_names=('imagenet2012',), **kw):
   """Returns config for discriminative_classifier of specified datasets."""
   config = ml_collections.ConfigDict(dict(
       dataset_names=list(dataset_names),
       type='proj.image_text.discriminative_classifier',
       prefix='z/0shot/',
       pp_img=_square875(sz),
       dataset_overrides={},
-      log_steps=log_steps,
       cache_final=True,
+      **kw,
   ))
   if pp_txt:
     config.pp_txt = pp_txt
@@ -100,18 +100,18 @@ def get_disclf(sz, *, log_steps, pp_txt=None, dataset_names=('imagenet2012',)):
 
 def get_coco(
     *,
-    log_steps,
     pp_img='resize(224)|value_range(-1, 1)',
     pp_txt='tokenize(max_len=16, inkey="texts", eos="sticky", pad_value=1)',
-    prefix='z/retr/coco_'):
+    prefix='z/retr/coco_',
+    **kw):
   """Returns config for mscoco retrieval zero-shot.
 
   Args:
-    log_steps: How often the evaluators should be run.
     pp_img: Pre-processing string for "image" feature.
     pp_txt: Pre-processing string for texts (expected to tokenize "texts" to
       "labels").
     prefix: Prefix to use for metrics.
+    **kw: Other config settings, most notably log_{steps,percent,...}.
 
   Returns:
     `ConfigDict` that can be used as a retrieval evaluator configuration.
@@ -123,5 +123,5 @@ def get_coco(
       'prefix': prefix,
       'dataset': 'coco_captions',
       'txt_name': ('captions', 'text'),
-      'log_steps': log_steps,
+      **kw,
   })

big_vision/configs/proj/image_text/lit.ipynb

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# Copyright 2023 Big Vision Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -129,14 +129,11 @@ def get_config(arg=None):
   ]
 
   config.grad_clip_norm = 1.0
-  # Gather representations across TPU cores for larger batch size for loss.
-  # See Figure 9 from https://arxiv.org/abs/2111.07991
-  config.loss_use_global_batch = True
 
   # Eval section (Both few-shot and zero-shot)
   eval_common = dict(
       type='proj.image_text.contrastive',
-      use_global_batch=config.loss_use_global_batch,
+      use_global_batch=True,
       log_steps=500 if not arg.runlocal else 5,
   )
   config.evals = {}

big_vision/configs/proj/image_text/lit_coco.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# Copyright 2023 Big Vision Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.