checkpoint.py

# Copyright 2024 Google LLC.
#
# 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.

import collections
from collections import abc
import re

from absl import logging
import flax
from flax.training import checkpoints
import jax.numpy as jnp
import numpy as np
from packaging import version
import pandas as pd
import scipy.ndimage
from tensorflow.io import gfile  # pylint: disable=import-error
import tqdm


def _flatten_dict(d, parent_key='', sep='/'):
  """Flattens a dictionary, keeping empty leaves."""
  items = []
  for k, v in d.items():
    path = parent_key + sep + k if parent_key else k
    if isinstance(v, abc.Mapping):
      items.extend(_flatten_dict(v, path, sep=sep).items())
    else:
      items.append((path, v))

  # Keeps the empty dict if it was set explicitly.
  if parent_key and not d:
    items.append((parent_key, {}))

  return dict(items)


def inspect_params(*,
                   params,
                   expected,
                   fail_if_extra=True,
                   fail_if_missing=True):
  """Inspects whether the params are consistent with the expected keys."""
  params_flat = _flatten_dict(params)
  expected_flat = _flatten_dict(expected)
  missing_keys = expected_flat.keys() - params_flat.keys()
  extra_keys = params_flat.keys() - expected_flat.keys()

  # Adds back empty dict explicitly, to support layers without weights.
  # Context: FLAX ignores empty dict during serialization.
  empty_keys = set()
  for k in missing_keys:
    if isinstance(expected_flat[k], dict) and not expected_flat[k]:
      params[k] = {}
      empty_keys.add(k)
  missing_keys -= empty_keys

  if empty_keys:
    logging.warning('Inspect recovered empty keys:\n%s', empty_keys)
  if missing_keys:
    logging.info('Inspect missing keys:\n%s', missing_keys)
  if extra_keys:
    logging.info('Inspect extra keys:\n%s', extra_keys)

  if (missing_keys and fail_if_missing) or (extra_keys and fail_if_extra):
    raise ValueError(f'Missing params from checkpoint: {missing_keys}.\n'
                     f'Extra params in checkpoint: {extra_keys}.\n'
                     f'Restored params from checkpoint: {params_flat.keys()}.\n'
                     f'Expected params from code: {expected_flat.keys()}.')
  return params


def recover_tree(keys, values):
  """Recovers a tree as a nested dict from flat names and values.

  This function is useful to analyze checkpoints that are without need to access
  the exact source code of the experiment. In particular, it can be used to
  extract an reuse various subtrees of the scheckpoint, e.g. subtree of
  parameters.

  Args:
    keys: a list of keys, where '/' is used as separator between nodes.
    values: a list of leaf values.

  Returns:
    A nested tree-like dict.
  """
  tree = {}
  sub_trees = collections.defaultdict(list)
  for k, v in zip(keys, values):
    if '/' not in k:
      tree[k] = v
    else:
      k_left, k_right = k.split('/', 1)
      sub_trees[k_left].append((k_right, v))
  for k, kv_pairs in sub_trees.items():
    k_subtree, v_subtree = zip(*kv_pairs)
    tree[k] = recover_tree(k_subtree, v_subtree)
  return tree


def copy(src, dst, progress=True, block_size=1024 * 1024 * 10):
  """Copies a file with progress bar.

  Args:
    src: Source file. Path must be readable by `tf.io.gfile`.
    dst: Destination file. Path must be readable by `tf.io.gfile`.
    progress: Whether to show a progres bar.
    block_size: Size of individual blocks to be read/written.
  """
  stats = gfile.stat(src)
  n = int(np.ceil(stats.length / block_size))
  range_or_trange = tqdm.trange if progress else range
  with gfile.GFile(src, 'rb') as fin:
    with gfile.GFile(dst, 'wb') as fout:
      for _ in range_or_trange(n):
        fout.write(fin.read(block_size))


def load(path):
  """Loads params from a checkpoint previously stored with `save()`."""
  with gfile.GFile(path, 'rb') as f:
    ckpt_dict = np.load(f, allow_pickle=False)
    keys, values = zip(*list(ckpt_dict.items()))
  params = checkpoints.convert_pre_linen(recover_tree(keys, values))
  if isinstance(params, flax.core.FrozenDict):
    params = params.unfreeze()
  if version.parse(flax.__version__) >= version.parse('0.3.6'):
    params = _fix_groupnorm(params)
  return params


def _fix_groupnorm(params):
  # See https://github.com/google/flax/issues/1721
  regex = re.compile(r'gn(\d+|_root|_proj)$')

  def fix_gn(args):
    path, array = args
    if len(path) > 1 and regex.match(
        path[-2]) and path[-1] in ('bias', 'scale'):
      array = array.squeeze()
    return (path, array)

  return flax.traverse_util.unflatten_dict(
      dict(map(fix_gn,
               flax.traverse_util.flatten_dict(params).items())))


def load_pretrained(*, pretrained_path, init_params, model_config):
  """加载/转换一个预训练的检查点进行微调

  Args:
    pretrained_path: File pointing to pretrained checkpoint.
    init_params: Parameters from model. Will be used for the head of the model
      and to verify that the model is compatible with the stored checkpoint.
    model_config: Configuration of the model. Will be used to configure the head
      and rescale the position embeddings.

  Returns:
    Parameters like `init_params`, but loaded with pretrained weights from
    `pretrained_path` and adapted accordingly.
  """

  restored_params = inspect_params(
      params=load(pretrained_path),
      expected=init_params,
      fail_if_extra=False,
      fail_if_missing=False)

  # The following allows implementing fine-tuning head variants depending on the
  # value of `representation_size` in the fine-tuning job:
  # - `None` : drop the whole head and attach a nn.Linear.
  # - same number as in pre-training means : keep the head but reset the last
  #    layer (logits) for the new task.
  if model_config.get('representation_size') is None:
    if 'pre_logits' in restored_params:
      logging.info('load_pretrained: drop-head variant')
      restored_params['pre_logits'] = {}
  restored_params['head']['kernel'] = init_params['head']['kernel']
  restored_params['head']['bias'] = init_params['head']['bias']

  if 'posembed_input' in restored_params.get('Transformer', {}):
    # Rescale the grid of position embeddings. Param shape is (1,N,1024)
    posemb = restored_params['Transformer']['posembed_input']['pos_embedding']
    posemb_new = init_params['Transformer']['posembed_input']['pos_embedding']
    if posemb.shape != posemb_new.shape:
      logging.info('load_pretrained: resized variant: %s to %s', posemb.shape,
                   posemb_new.shape)
      posemb = interpolate_posembed(
          posemb, posemb_new.shape[1], model_config.classifier == 'token')
      restored_params['Transformer']['posembed_input']['pos_embedding'] = posemb

  if version.parse(flax.__version__) >= version.parse('0.3.6'):
    restored_params = _fix_groupnorm(restored_params)

  return flax.core.freeze(restored_params)


def interpolate_posembed(posemb, num_tokens: int, has_class_token: bool):
  """Interpolate given positional embedding parameters into a new shape.

  Args:
    posemb: positional embedding parameters.
    num_tokens: desired number of tokens.
    has_class_token: True if the positional embedding parameters contain a
      class token.

  Returns:
    Positional embedding parameters interpolated into the new shape.
  """
  assert posemb.shape[0] == 1
  if has_class_token:
    posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
    num_tokens -= 1
  else:
    posemb_tok, posemb_grid = posemb[:, :0], posemb[0, 0:]

  gs_old = int(np.sqrt(len(posemb_grid)))
  gs_new = int(np.sqrt(num_tokens))
  logging.info('interpolate_posembed: grid-size from %s to %s', gs_old, gs_new)
  assert gs_old ** 2 == len(posemb_grid), f'{gs_old ** 2} != {len(posemb_grid)}'
  assert gs_new ** 2 == num_tokens, f'{gs_new ** 2} != {num_tokens}'
  posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)

  zoom = (gs_new / gs_old, gs_new / gs_old, 1)
  posemb_grid = scipy.ndimage.zoom(posemb_grid, zoom, order=1)
  posemb_grid = posemb_grid.reshape(1, gs_new * gs_new, -1)
  return jnp.array(np.concatenate([posemb_tok, posemb_grid], axis=1))


def get_augreg_df(directory='gs://vit_models/augreg'):
  """Reads DataFrame describing AugReg models from GCS bucket.

  This function returns a dataframe that describes the models that were
  published as part of the paper "How to train your ViT? Data, Augmentation, and
  Regularization in Vision Transformers" (https://arxiv.org/abs/TODO).

  Note that every row in the dataset corresponds to a pre-training checkpoint
  (column "filename"), and a fine-tuning checkpoint (column "adapt_filename").
  Every pre-trained checkpoint is fine-tuned many times.

  Args:
    directory: Pathname of directory containing "index.csv"

  Returns:
    Dataframe with the following columns:
      - name: Name of the model, as used in descriptions in paper (e.g. "B/16",
        or "R26+S/32").
      - ds: Dataset used for pre-training: "i1k" (300 epochs), "i21k" (300
        epochs), and "i21k_30" (30 epochs).
      - lr: Learning rate used for pre-training.
      - aug: Data augmentation used for pre-training. Refer to paper for
        details.
      - wd: Weight decay used for pre-training.
      - do: Dropout used for pre-training.
      - sd: Stochastic depth used for pre-training.
      - best_val: Best accuracy on validation set that was reached during the
        pre-training. Note that "validation set" can refer to minival (meaning
        split from training set, as for example for "imagenet2012" dataset).
      - final_val: Final validation set accuracy.
      - final_test: Final testset accuracy (in cases where there is no official
        testset, like for "imagenet2012", this refers to the validation set).
      - adapt_ds: What dataset was used for fine-tuning.
      - adapt_lr: Learning rate used for fine-tuning.
      - adapt_steps: Number of steps used for fine-tuning (with a fixed batch
        size of 512).
      - adapt_resolution: Resolution that was used for fine-tuning.
      - adapt_final_val: Final validation accuracy after fine-tuning.
      - adapt_final_test: Final test accuracy after fine-tuning.
      - params: Number of parameters.
      - infer_samples_per_sec: Numbers of sample per seconds during inference on
        a V100 GPU (measured with `timm` implementation).
      - filename: Name of the pre-training checkpoint. Can be found at
        "gs://vit_models/augreg/{filename}.npz".
      - adapt_filename: Name of the fine-tuning checkpoint.
  """
  with gfile.GFile(f'{directory}/index.csv') as f:
    return pd.read_csv(f)