Load ViT with CLIPPO Weights about big_vision

[ahmadmustafaanis]

I also tried with the load function in vit.

from big_vision.models.vit import load as vit_load
init_params = vit_load(None, 'clippo_b16_yfcc100m_i21k_init_75c4.npz', None)

and in vit_load

def fix_old_checkpoints(params):
    """Fix small bwd incompat that can't be resolved with names in model def."""

    params = flax.core.unfreeze(flax.training.checkpoints.convert_pre_linen(params))
    # Original ViT paper variant had posemb in a module:
    if "posembed_input" in params["img"]:  # IT WAS params['Transformer'] before which I replaced as no transformer key in weights
        logging.info("ViT: Loading and fixing VERY old posemb")
        posemb = params["img"].pop("posembed_input")
        params["pos_embedding"] = posemb["pos_embedding"]

    # Widely used version before 2022 had posemb in Encoder:
    if "pos_embedding" in params["img"]:
        logging.info("ViT: Loading and fixing old posemb")
        params["pos_embedding"] = params["img"].pop("pos_embedding")

    # Old vit.py used to first concat [cls] token, then add posemb.
    # This means a B/32@224px would have 7x7+1 posembs. This is useless and clumsy
    # so we changed to add posemb then concat [cls]. We can recover the old
    # checkpoint by manually summing [cls] token and its posemb entry.
    if "pos_embedding" in params:
        pe = params["pos_embedding"]
        if int(np.sqrt(pe.shape[1])) ** 2 + 1 == int(pe.shape[1]):
            logging.info("ViT: Loading and fixing combined cls+posemb")
            pe_cls, params["pos_embedding"] = pe[:, :1], pe[:, 1:]
            if "cls" in params:
                params["cls"] += pe_cls

    # MAP-head variants during ViT-G development had it inlined:
    if "probe" in params:
        params["MAPHead_0"] = {
            k: params.pop(k)
            for k in [
                "probe",
                "MlpBlock_0",
                "MultiHeadDotProductAttention_0",
                "LayerNorm_0",
            ]
        }

    return params

def load(
    init_params, init_file, model_cfg, dont_load=()
):  # pylint: disable=invalid-name because we had to CamelCase above.
    """Load init from checkpoint, both old model and this one. +Hi-res posemb."""
    del model_cfg

    init_file = VANITY_NAMES.get(init_file, init_file)
    restored_params = utils.load_params(None, init_file)

    restored_params = fix_old_checkpoints(restored_params)

    # possibly use the random init for some of the params (such as, the head).
    restored_params = common.merge_params(restored_params, init_params, dont_load)

    # resample posemb if needed.
    if init_params and "pos_embedding" in init_params:
        restored_params["pos_embedding"] = resample_posemb(
            old=restored_params["pos_embedding"], new=init_params["pos_embedding"]
        )

    return restored_params

I still get the same error:

---------------------------------------------------------------------------
ScopeParamNotFoundError                   Traceback (most recent call last)
Cell In[10], [line 48](vscode-notebook-cell:?execution_count=10&line=48)
     [45](vscode-notebook-cell:?execution_count=10&line=45) labels = batch['label'].numpy()
     [47](vscode-notebook-cell:?execution_count=10&line=47) # Forward pass
---> [48](vscode-notebook-cell:?execution_count=10&line=48) zimg, _ = model.apply({'params': init_params}, images)  
     [50](vscode-notebook-cell:?execution_count=10&line=50) # Get logits
     [51](vscode-notebook-cell:?execution_count=10&line=51) logits = classifier.apply({'params': classifier.params}, zimg)

    [... skipping hidden 6 frame]

File [~/Desktop/projects/big_vision/big_vision/models/vit.py:186](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:186), in _Model.__call__(self, image, train)
    [183](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:183) out = {}
    [185](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:185) # Patch extraction
--> [186](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:186) x = out["stem"] = nn.Conv(
    [187](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:187)     self.width,
    [188](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:188)     self.patch_size,
    [189](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:189)     strides=self.patch_size,
    [190](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:190)     padding="VALID",
    [191](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:191)     name="embedding",
    [192](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:192) )(image)
    [194](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:194) n, h, w, c = x.shape
    [195](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/Desktop/projects/big_vision/big_vision/models/vit.py:195) x = jnp.reshape(x, [n, h * w, c])

    [... skipping hidden 2 frame]

File [~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:480](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:480), in _Conv.__call__(self, inputs)
    [474](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:474) if self.mask is not None and self.mask.shape != kernel_shape:
    [475](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:475)   raise ValueError(
    [476](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:476)       'Mask needs to have the same shape as weights. '
    [477](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:477)       f'Shapes are: {self.mask.shape}, {kernel_shape}'
    [478](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:478)   )
--> [480](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:480) kernel = self.param(
    [481](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:481)     'kernel', self.kernel_init, kernel_shape, self.param_dtype
    [482](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:482) )
    [484](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:484) if self.mask is not None:
    [485](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/linen/linear.py:485)   kernel *= self.mask

    [... skipping hidden 1 frame]

File [~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:896](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:896), in Scope.param(self, name, init_fn, unbox, *init_args)
    [894](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:894)   if self.is_collection_empty('params'):
    [895](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:895)     raise errors.ScopeCollectionNotFound('params', name, self.path_text)
--> [896](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:896)   raise errors.ScopeParamNotFoundError(name, self.path_text)
    [897](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:897) value = init_fn(self.make_rng('params'), *init_args)
    [898](https://file+.vscode-resource.vscode-cdn.net/home/ahmad/Desktop/projects/big_vision/big_vision/trainers/~/anaconda3/envs/robustness/lib/python3.8/site-packages/flax/core/scope.py:898) self.put_variable('params', name, value)

ScopeParamNotFoundError: Could not find parameter named "kernel" in scope "/embedding". (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamNotFoundError)

How ever, there is now an extra key(pos_embedding) in init_params with this dict_keys(['img', 't', 'pos_embedding']) as we created it in the function

[mitscha]

Hi,

I think the issue you are observing is because CLIPPO uses a wrapper for models.vit (to do two forward passes with natural and text images), namely models.proj.clippo.one_tower. Loading the checkpoints into this model should work.

You can follow the steps in the colab. This might be broken at head, but using your method to check out the older commit (or the follow up commit which adds the colab) should work.