<jupyter_start><jupyter_text>See code at https://github.com/google-research/vision_transformer/See papers at- Vision Transformer: https://arxiv.org/abs/2010.11929- MLP-Mixer: https://arxiv.org/abs/2105.01601- How to train your ViT: https://arxiv.org/abs/2106.10270- When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations: https://arxiv.org/abs/2106.01548This Colab allows you to run the [JAX](https://jax.readthedocs.org) implementation of the Vision Transformer.If you just want to load a pre-trained checkpoint from a large repository anddirectly use it for inference, you probably want to go the other Colabhttps://colab.sandbox.google.com/github/google-research/vision_transformer/blob/linen/vit_jax_augreg.ipynb<jupyter_code>import os
from google.colab import drive

drive.mount('/content/drive',force_remount=True)
os.chdir("/content/drive/MyDrive/Colab Notebooks")
%%bash
rm -rf vision_transformer
git clone https://github.com/missingstuffedbun/vision_transformer.git
pip install -qr vision_transformer/vit_jax/requirements.txt
os.chdir("/content/drive/MyDrive/Colab Notebooks/vision_transformer")

# %%bash
# mkdir -p /content/input
# mkdir -p /content/working
# mkdir -p /content/input/sarship
# mkdir -p /content/input/sarship/train
# mkdir -p /content/input/sarship/val
# mkdir -p /content/input/sarship/test
# tar -C /content/input/sarship/train -xf ./sarship/train.tar
# tar -C /content/input/sarship/val -xf ./sarship/val.tar
# tar -C /content/input/sarship/test -xf ./sarship/test.tar
# %%bash
# python detr/create_coco_tf_record.py --logtostderr \
# --train_image_dir=$DATAPATH/input/sarship/train/train \
# --val_image_dir=$DATAPATH/input/sarship/val/val \
# --test_image_dir=$DATAPATH/input/sarship/test/test \
# --train_annotations_file=$DATAPATH/working/custom_train.json \
# --val_annotations_file=$DATAPATH/working/custom_val.json \
# --test_annotations_file=$DATAPATH/working/custom_test.json \
# --output_dir=$DATAPATH/input/tf<jupyter_output><empty_output><jupyter_text>Load dataset<jupyter_code>import sys
if './vision_transformer' not in sys.path:
  sys.path.append('./vision_transformer')

from vit_jax import input_pipeline

# For details about setting up datasets, see input_pipeline.py on the right.

ds_train = input_pipeline.get_data_from_tfds(config=config, mode='train')
ds_test = input_pipeline.get_data_from_tfds(config=config, mode='test')

del config  # Only needed to instantiate datasets.
# Fetch a batch of test images for illustration purposes.
batch = next(iter(ds_test.as_numpy_iterator()))
# Note the shape : [num_local_devices, local_batch_size, h, w, c]
batch['image'].shape
# Show some imags with their labels.
images, labels = batch['image'][0][:9], batch['label'][0][:9]
titles = map(make_label_getter(dataset), labels.argmax(axis=1))
show_img_grid(images, titles)
# Same as above, but with train images.
# Note how images are cropped/scaled differently.
# Check out input_pipeline.get_data() in the editor at your right to see how the
# images are preprocessed differently.
batch = next(iter(ds_train.as_numpy_iterator()))
images, labels = batch['image'][0][:9], batch['label'][0][:9]
titles = map(make_label_getter(dataset), labels.argmax(axis=1))
show_img_grid(images, titles)<jupyter_output><empty_output><jupyter_text>Load pre-trained<jupyter_code>model_config = models_config.MODEL_CONFIGS[model_name]
model_config
# Load model definition & initialize random parameters.
# This also compiles the model to XLA (takes some minutes the first time).
if model_name.startswith('Mixer'):
  model = models.MlpMixer(num_classes=num_classes, **model_config)
else:
  model = models.VisionTransformer(num_classes=num_classes, **model_config)
variables = jax.jit(lambda: model.init(
    jax.random.PRNGKey(0),
    # Discard the "num_local_devices" dimension of the batch for initialization.
    batch['image'][0, :1],
    train=False,
), backend='cpu')()
# Load and convert pretrained checkpoint.
# This involves loading the actual pre-trained model results, but then also also
# modifying the parameters a bit, e.g. changing the final layers, and resizing
# the positional embeddings.
# For details, refer to the code and to the methods of the paper.
params = checkpoint.load_pretrained(
    pretrained_path=f'{model_name}.npz',
    init_params=variables['params'],
    model_config=model_config,
)<jupyter_output>INFO:absl:Inspect extra keys:
{'pre_logits/bias', 'pre_logits/kernel'}
INFO:absl:load_pretrained: drop-head variant<jupyter_text>Evaluate<jupyter_code># So far, all our data is in the host memory. Let's now replicate the arrays
# into the devices.
# This will make every array in the pytree params become a ShardedDeviceArray
# that has the same data replicated across all local devices.
# For TPU it replicates the params in every core.
# For a single GPU this simply moves the data onto the device.
# For CPU it simply creates a copy.
params_repl = flax.jax_utils.replicate(params)
print('params.cls:', type(params['head']['bias']).__name__,
      params['head']['bias'].shape)
print('params_repl.cls:', type(params_repl['head']['bias']).__name__,
      params_repl['head']['bias'].shape)
# Then map the call to our model's forward pass onto all available devices.
vit_apply_repl = jax.pmap(lambda params, inputs: model.apply(
    dict(params=params), inputs, train=False))
def get_accuracy(params_repl):
  """Returns accuracy evaluated on the test set."""
  good = total = 0
  steps = input_pipeline.get_dataset_info(dataset, 'test')['num_examples'] // batch_size
  for _, batch in zip(tqdm.trange(steps), ds_test.as_numpy_iterator()):
    predicted = vit_apply_repl(params_repl, batch['image'])
    is_same = predicted.argmax(axis=-1) == batch['label'].argmax(axis=-1)
    good += is_same.sum()
    total += len(is_same.flatten())
  return good / total
# Random performance without fine-tuning.
get_accuracy(params_repl)<jupyter_output>INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
100%|██████████| 19/19 [01:07<00:00,  3.58s/it]<jupyter_text>Fine-tune<jupyter_code># 100 Steps take approximately 15 minutes in the TPU runtime.
total_steps = 100
warmup_steps = 5
decay_type = 'cosine'
grad_norm_clip = 1
# This controls in how many forward passes the batch is split. 8 works well with
# a TPU runtime that has 8 devices. 64 should work on a GPU. You can of course
# also adjust the batch_size above, but that would require you to adjust the
# learning rate accordingly.
accum_steps = 8
base_lr = 0.03
# Check out train.make_update_fn in the editor on the right side for details.
lr_fn = utils.create_learning_rate_schedule(total_steps, base_lr, decay_type, warmup_steps)
update_fn_repl = train.make_update_fn(
    apply_fn=model.apply, accum_steps=accum_steps, lr_fn=lr_fn)
# We use a momentum optimizer that uses half precision for state to save
# memory. It als implements the gradient clipping.
opt = momentum_clip.Optimizer(grad_norm_clip=grad_norm_clip).create(params)
opt_repl = flax.jax_utils.replicate(opt)
# Initialize PRNGs for dropout.
update_rng_repl = flax.jax_utils.replicate(jax.random.PRNGKey(0))
losses = []
lrs = []
# Completes in ~20 min on the TPU runtime.
for step, batch in zip(
    tqdm.trange(1, total_steps + 1),
    ds_train.as_numpy_iterator(),
):

  opt_repl, loss_repl, update_rng_repl = update_fn_repl(
      opt_repl, flax.jax_utils.replicate(step), batch, update_rng_repl)
  losses.append(loss_repl[0])
  lrs.append(lr_fn(step))

plt.plot(losses)
plt.figure()
plt.plot(lrs)

# Should be ~96.7% for Mixer-B/16 or 97.7% for ViT-B/32 on CIFAR10 (both @224)
get_accuracy(opt_repl.target)<jupyter_output>INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
100%|██████████| 19/19 [00:32<00:00,  1.73s/it]<jupyter_text>Inference<jupyter_code># Download a pre-trained model.

if model_name.startswith('Mixer'):
  # Download model trained on imagenet2012
  ![ -e "$model_name"_imagenet2012.npz ] || gsutil cp gs://mixer_models/imagenet1k/"$model_name".npz "$model_name"_imagenet2012.npz
  model = models.MlpMixer(num_classes=1000, **model_config)
else:
  # Download model pre-trained on imagenet21k and fine-tuned on imagenet2012.
  ![ -e "$model_name"_imagenet2012.npz ] || gsutil cp gs://vit_models/imagenet21k+imagenet2012/"$model_name".npz "$model_name"_imagenet2012.npz
  model = models.VisionTransformer(num_classes=1000, **model_config)

import os
assert os.path.exists(f'{model_name}_imagenet2012.npz')
# Load and convert pretrained checkpoint.
params = checkpoint.load(f'{model_name}_imagenet2012.npz')
params['pre_logits'] = {}  # Need to restore empty leaf for Flax.
# Get imagenet labels.
!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt
imagenet_labels = dict(enumerate(open('ilsvrc2012_wordnet_lemmas.txt')))
# Get a random picture with the correct dimensions.
resolution = 224 if model_name.startswith('Mixer') else 384
!wget https://picsum.photos/$resolution -O picsum.jpg
import PIL
img = PIL.Image.open('picsum.jpg')
img
# Predict on a batch with a single item (note very efficient TPU usage...)
logits, = model.apply(dict(params=params), (np.array(img) / 128 - 1)[None, ...], train=False)
preds = flax.nn.softmax(logits)
for idx in preds.argsort()[:-11:-1]:
  print(f'{preds[idx]:.5f} : {imagenet_labels[idx]}', end='')<jupyter_output>0.13330 : sandbar, sand_bar
0.09332 : seashore, coast, seacoast, sea-coast
0.05257 : jeep, landrover
0.05188 : Arabian_camel, dromedary, Camelus_dromedarius
0.01251 : horned_viper, cerastes, sand_viper, horned_asp, Cerastes_cornutus
0.00753 : tiger_beetle
0.00744 : dung_beetle
0.00711 : sidewinder, horned_rattlesnake, Crotalus_cerastes
0.00703 : leatherback_turtle, leatherback, leathery_turtle, Dermochelys_coriacea
0.00647 : pole
{'fraction_non_alphanumeric': 0.12178177360032694, 'fraction_numerical': 0.027380465876583573, 'mean_word_length': 5.552208835341365, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 21, 'https://': 10, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '597535', 'license': 'Apache-2.0', 'n_tokens_mistral': 3627, 'n_tokens_neox': 3391, 'n_words': 821, 'path': 'VIT_WWZ.ipynb', 'repo_name': 'missingstuffedbun/vision_transformer', 'size': '9788'}