search

google-research / big_vision

Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.
Apache License 2.0
2.25k stars 147 forks source link

Load ViT with CLIPPO Weights #88

Closed ahmadmustafaanis closed 8 months ago

ahmadmustafaanis commented 8 months ago

Hi, I am trying to load the ViT model in big vision with the clippo weights.

!git clone --branch=main https://github.com/google-research/big_vision
!cd big_vision && git checkout fd2d3bd2efc9d89ea959f16cd2f58ae8a495cd44  # this is the clippo commit

now my script:

checkpoint_path = '/home/ahmad/Desktop/projects/big_vision/big_vision/weights/clippo_b16_yfcc100m_i21k_init_75c4.npz'
init_params = utils.load_checkpoint(None, checkpoint_path)['params']
model = vit.Model(
  num_classes=768, variant=None, name="img"
)
classifier = nn.Dense(10) # 10 classes
opt = optax.adam(1e-3)
# Training 
for epoch in range(10):

  for batch in ds['train'].batch(2):

    # Get images and labels
    images = batch['image'].numpy()
    # convert the image from (28, 28, 1) gray scale to (28, 28, 3) rgb
    images = np.repeat(images, 3, axis=-1)
    print(images.shape)
    labels = batch['label'].numpy()

    # Forward pass
    zimg, _ = model.apply({'params': init_params}, images)

This is the error I get

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

The npz file has the following keys:

dict_keys(['chrono', 'opt', 'params'])

and in the params parameter keys are:

dict_keys(['t', 'img'])

and in the img we have

init_params['params']['img'].keys()
>>>dict_keys(['pos_embedding', 'MAPHead_0', 'Transformer', 'embedding', 'head'])

My Jax version is:

'0.4.13'

and flax version is:

'0.7.2'
ahmadmustafaanis commented 8 months ago

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 commented 8 months ago

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.