match_steam_banners_with_BEiT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "match_steam_banners_with_BEiT.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "IdZc-08r0WRs"
      },
      "source": [
        "%pip install -q transformers mediapy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HdWexgxf0hrw"
      },
      "source": [
        "import torch\n",
        "import mediapy as media\n",
        "\n",
        "from transformers import BeitFeatureExtractor\n",
        "from transformers import BeitForMaskedImageModeling, BeitForImageClassification\n",
        "from PIL import Image\n",
        "import requests\n",
        "\n",
        "url = 'http://images.cocodataset.org/val2017/000000039769.jpg'\n",
        "image = Image.open(requests.get(url, stream=True).raw)\n",
        "\n",
        "# TODO: base vs. large\n",
        "keyword = 'base'\n",
        "# TODO: simple vs. complex\n",
        "feature_style = 'complex'\n",
        "\n",
        "model_name = f'microsoft/beit-{keyword}-patch16-224-pt22k'\n",
        "# model_name = f'microsoft/beit-{keyword}-patch16-224-pt22k-ft22k'\n",
        "# model_name = f'microsoft/beit-{keyword}-patch16-224'\n",
        "\n",
        "is_finetuned = not model_name.endswith('pt22k')\n",
        "\n",
        "if is_finetuned:\n",
        "  arch = BeitForImageClassification\n",
        "else:\n",
        "  arch = BeitForMaskedImageModeling\n",
        "\n",
        "feature_extractor = BeitFeatureExtractor.from_pretrained(model_name)\n",
        "model = arch.from_pretrained(model_name)\n",
        "\n",
        "# Delete the classification head\n",
        "if is_finetuned:\n",
        "  model.classifier = torch.nn.Identity()\n",
        "else:\n",
        "  model.lm_head = torch.nn.Identity()  \n",
        "  # HuggingFace suggests the classification head may include `model.layernorm`.\n",
        "  # However, I do not follow this view and avoid setting this layer to Identity.\n",
        "\n",
        "inputs = feature_extractor(images=image, return_tensors=\"pt\")\n",
        "outputs = model(**inputs)\n",
        "\n",
        "if is_finetuned:\n",
        "  pooled_output = outputs.logits\n",
        "else:\n",
        "  hidden_states = outputs.logits\n",
        "\n",
        "  # The final hidden state of the [CLS] token\n",
        "  cls_token = hidden_states[:, 0]\n",
        "  # The final hidden states of the patch tokens\n",
        "  patch_tokens = hidden_states[:, 1:, :]\n",
        "\n",
        "  if feature_style == 'simple':\n",
        "    pooled_output = cls_token\n",
        "  else:\n",
        "    # Mean pool\n",
        "    # https://github.com/huggingface/transformers/blob/bcc3f7b6560c1ed427f051107c7755956a27a9f2/src/transformers/models/beit/modeling_beit.py#L664-L670\n",
        "    mean_token = patch_tokens.mean(1)\n",
        "    # Concatenate, as with ViT-Base for the eval_linear of DINO\n",
        "    # https://github.com/facebookresearch/dino/blob/499d9e2b3d903355f67e86f9def06bccb4222b1f/eval_linear.py#L256-L260\n",
        "    pooled_output = torch.cat([cls_token, mean_token], dim=-1)\n",
        "\n",
        "print(outputs.keys())\n",
        "for i in outputs.values():\n",
        "  print(i.size())\n",
        "\n",
        "if not is_finetuned:\n",
        "  print(pooled_output.size())"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}