convert_beit_ckpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "module.positional_embedding\n",
      "module.text_projection\n",
      "module.logit_scale\n",
      "module.visual.class_embedding\n",
      "module.visual.positional_embedding\n",
      "module.visual.proj\n",
      "module.visual.conv1.weight\n",
      "module.visual.ln_pre.weight\n",
      "module.visual.ln_pre.bias\n",
      "module.visual.transformer.resblocks.0.ln_1.weight\n",
      "module.visual.transformer.resblocks.0.ln_1.bias\n",
      "module.visual.transformer.resblocks.0.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.0.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.0.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.0.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.0.ln_2.weight\n",
      "module.visual.transformer.resblocks.0.ln_2.bias\n",
      "module.visual.transformer.resblocks.0.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.0.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.0.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.0.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.1.ln_1.weight\n",
      "module.visual.transformer.resblocks.1.ln_1.bias\n",
      "module.visual.transformer.resblocks.1.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.1.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.1.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.1.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.1.ln_2.weight\n",
      "module.visual.transformer.resblocks.1.ln_2.bias\n",
      "module.visual.transformer.resblocks.1.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.1.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.1.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.1.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.2.ln_1.weight\n",
      "module.visual.transformer.resblocks.2.ln_1.bias\n",
      "module.visual.transformer.resblocks.2.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.2.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.2.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.2.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.2.ln_2.weight\n",
      "module.visual.transformer.resblocks.2.ln_2.bias\n",
      "module.visual.transformer.resblocks.2.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.2.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.2.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.2.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.3.ln_1.weight\n",
      "module.visual.transformer.resblocks.3.ln_1.bias\n",
      "module.visual.transformer.resblocks.3.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.3.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.3.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.3.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.3.ln_2.weight\n",
      "module.visual.transformer.resblocks.3.ln_2.bias\n",
      "module.visual.transformer.resblocks.3.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.3.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.3.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.3.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.4.ln_1.weight\n",
      "module.visual.transformer.resblocks.4.ln_1.bias\n",
      "module.visual.transformer.resblocks.4.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.4.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.4.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.4.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.4.ln_2.weight\n",
      "module.visual.transformer.resblocks.4.ln_2.bias\n",
      "module.visual.transformer.resblocks.4.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.4.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.4.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.4.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.5.ln_1.weight\n",
      "module.visual.transformer.resblocks.5.ln_1.bias\n",
      "module.visual.transformer.resblocks.5.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.5.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.5.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.5.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.5.ln_2.weight\n",
      "module.visual.transformer.resblocks.5.ln_2.bias\n",
      "module.visual.transformer.resblocks.5.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.5.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.5.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.5.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.6.ln_1.weight\n",
      "module.visual.transformer.resblocks.6.ln_1.bias\n",
      "module.visual.transformer.resblocks.6.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.6.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.6.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.6.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.6.ln_2.weight\n",
      "module.visual.transformer.resblocks.6.ln_2.bias\n",
      "module.visual.transformer.resblocks.6.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.6.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.6.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.6.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.7.ln_1.weight\n",
      "module.visual.transformer.resblocks.7.ln_1.bias\n",
      "module.visual.transformer.resblocks.7.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.7.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.7.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.7.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.7.ln_2.weight\n",
      "module.visual.transformer.resblocks.7.ln_2.bias\n",
      "module.visual.transformer.resblocks.7.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.7.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.7.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.7.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.8.ln_1.weight\n",
      "module.visual.transformer.resblocks.8.ln_1.bias\n",
      "module.visual.transformer.resblocks.8.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.8.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.8.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.8.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.8.ln_2.weight\n",
      "module.visual.transformer.resblocks.8.ln_2.bias\n",
      "module.visual.transformer.resblocks.8.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.8.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.8.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.8.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.9.ln_1.weight\n",
      "module.visual.transformer.resblocks.9.ln_1.bias\n",
      "module.visual.transformer.resblocks.9.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.9.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.9.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.9.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.9.ln_2.weight\n",
      "module.visual.transformer.resblocks.9.ln_2.bias\n",
      "module.visual.transformer.resblocks.9.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.9.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.9.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.9.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.10.ln_1.weight\n",
      "module.visual.transformer.resblocks.10.ln_1.bias\n",
      "module.visual.transformer.resblocks.10.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.10.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.10.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.10.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.10.ln_2.weight\n",
      "module.visual.transformer.resblocks.10.ln_2.bias\n",
      "module.visual.transformer.resblocks.10.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.10.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.10.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.10.mlp.c_proj.bias\n",
      "module.visual.transformer.resblocks.11.ln_1.weight\n",
      "module.visual.transformer.resblocks.11.ln_1.bias\n",
      "module.visual.transformer.resblocks.11.attn.in_proj_weight\n",
      "module.visual.transformer.resblocks.11.attn.in_proj_bias\n",
      "module.visual.transformer.resblocks.11.attn.out_proj.weight\n",
      "module.visual.transformer.resblocks.11.attn.out_proj.bias\n",
      "module.visual.transformer.resblocks.11.ln_2.weight\n",
      "module.visual.transformer.resblocks.11.ln_2.bias\n",
      "module.visual.transformer.resblocks.11.mlp.c_fc.weight\n",
      "module.visual.transformer.resblocks.11.mlp.c_fc.bias\n",
      "module.visual.transformer.resblocks.11.mlp.c_proj.weight\n",
      "module.visual.transformer.resblocks.11.mlp.c_proj.bias\n",
      "module.visual.ln_post.weight\n",
      "module.visual.ln_post.bias\n",
      "module.transformer.resblocks.0.ln_1.weight\n",
      "module.transformer.resblocks.0.ln_1.bias\n",
      "module.transformer.resblocks.0.attn.in_proj_weight\n",
      "module.transformer.resblocks.0.attn.in_proj_bias\n",
      "module.transformer.resblocks.0.attn.out_proj.weight\n",
      "module.transformer.resblocks.0.attn.out_proj.bias\n",
      "module.transformer.resblocks.0.ln_2.weight\n",
      "module.transformer.resblocks.0.ln_2.bias\n",
      "module.transformer.resblocks.0.mlp.c_fc.weight\n",
      "module.transformer.resblocks.0.mlp.c_fc.bias\n",
      "module.transformer.resblocks.0.mlp.c_proj.weight\n",
      "module.transformer.resblocks.0.mlp.c_proj.bias\n",
      "module.transformer.resblocks.1.ln_1.weight\n",
      "module.transformer.resblocks.1.ln_1.bias\n",
      "module.transformer.resblocks.1.attn.in_proj_weight\n",
      "module.transformer.resblocks.1.attn.in_proj_bias\n",
      "module.transformer.resblocks.1.attn.out_proj.weight\n",
      "module.transformer.resblocks.1.attn.out_proj.bias\n",
      "module.transformer.resblocks.1.ln_2.weight\n",
      "module.transformer.resblocks.1.ln_2.bias\n",
      "module.transformer.resblocks.1.mlp.c_fc.weight\n",
      "module.transformer.resblocks.1.mlp.c_fc.bias\n",
      "module.transformer.resblocks.1.mlp.c_proj.weight\n",
      "module.transformer.resblocks.1.mlp.c_proj.bias\n",
      "module.transformer.resblocks.2.ln_1.weight\n",
      "module.transformer.resblocks.2.ln_1.bias\n",
      "module.transformer.resblocks.2.attn.in_proj_weight\n",
      "module.transformer.resblocks.2.attn.in_proj_bias\n",
      "module.transformer.resblocks.2.attn.out_proj.weight\n",
      "module.transformer.resblocks.2.attn.out_proj.bias\n",
      "module.transformer.resblocks.2.ln_2.weight\n",
      "module.transformer.resblocks.2.ln_2.bias\n",
      "module.transformer.resblocks.2.mlp.c_fc.weight\n",
      "module.transformer.resblocks.2.mlp.c_fc.bias\n",
      "module.transformer.resblocks.2.mlp.c_proj.weight\n",
      "module.transformer.resblocks.2.mlp.c_proj.bias\n",
      "module.transformer.resblocks.3.ln_1.weight\n",
      "module.transformer.resblocks.3.ln_1.bias\n",
      "module.transformer.resblocks.3.attn.in_proj_weight\n",
      "module.transformer.resblocks.3.attn.in_proj_bias\n",
      "module.transformer.resblocks.3.attn.out_proj.weight\n",
      "module.transformer.resblocks.3.attn.out_proj.bias\n",
      "module.transformer.resblocks.3.ln_2.weight\n",
      "module.transformer.resblocks.3.ln_2.bias\n",
      "module.transformer.resblocks.3.mlp.c_fc.weight\n",
      "module.transformer.resblocks.3.mlp.c_fc.bias\n",
      "module.transformer.resblocks.3.mlp.c_proj.weight\n",
      "module.transformer.resblocks.3.mlp.c_proj.bias\n",
      "module.transformer.resblocks.4.ln_1.weight\n",
      "module.transformer.resblocks.4.ln_1.bias\n",
      "module.transformer.resblocks.4.attn.in_proj_weight\n",
      "module.transformer.resblocks.4.attn.in_proj_bias\n",
      "module.transformer.resblocks.4.attn.out_proj.weight\n",
      "module.transformer.resblocks.4.attn.out_proj.bias\n",
      "module.transformer.resblocks.4.ln_2.weight\n",
      "module.transformer.resblocks.4.ln_2.bias\n",
      "module.transformer.resblocks.4.mlp.c_fc.weight\n",
      "module.transformer.resblocks.4.mlp.c_fc.bias\n",
      "module.transformer.resblocks.4.mlp.c_proj.weight\n",
      "module.transformer.resblocks.4.mlp.c_proj.bias\n",
      "module.transformer.resblocks.5.ln_1.weight\n",
      "module.transformer.resblocks.5.ln_1.bias\n",
      "module.transformer.resblocks.5.attn.in_proj_weight\n",
      "module.transformer.resblocks.5.attn.in_proj_bias\n",
      "module.transformer.resblocks.5.attn.out_proj.weight\n",
      "module.transformer.resblocks.5.attn.out_proj.bias\n",
      "module.transformer.resblocks.5.ln_2.weight\n",
      "module.transformer.resblocks.5.ln_2.bias\n",
      "module.transformer.resblocks.5.mlp.c_fc.weight\n",
      "module.transformer.resblocks.5.mlp.c_fc.bias\n",
      "module.transformer.resblocks.5.mlp.c_proj.weight\n",
      "module.transformer.resblocks.5.mlp.c_proj.bias\n",
      "module.transformer.resblocks.6.ln_1.weight\n",
      "module.transformer.resblocks.6.ln_1.bias\n",
      "module.transformer.resblocks.6.attn.in_proj_weight\n",
      "module.transformer.resblocks.6.attn.in_proj_bias\n",
      "module.transformer.resblocks.6.attn.out_proj.weight\n",
      "module.transformer.resblocks.6.attn.out_proj.bias\n",
      "module.transformer.resblocks.6.ln_2.weight\n",
      "module.transformer.resblocks.6.ln_2.bias\n",
      "module.transformer.resblocks.6.mlp.c_fc.weight\n",
      "module.transformer.resblocks.6.mlp.c_fc.bias\n",
      "module.transformer.resblocks.6.mlp.c_proj.weight\n",
      "module.transformer.resblocks.6.mlp.c_proj.bias\n",
      "module.transformer.resblocks.7.ln_1.weight\n",
      "module.transformer.resblocks.7.ln_1.bias\n",
      "module.transformer.resblocks.7.attn.in_proj_weight\n",
      "module.transformer.resblocks.7.attn.in_proj_bias\n",
      "module.transformer.resblocks.7.attn.out_proj.weight\n",
      "module.transformer.resblocks.7.attn.out_proj.bias\n",
      "module.transformer.resblocks.7.ln_2.weight\n",
      "module.transformer.resblocks.7.ln_2.bias\n",
      "module.transformer.resblocks.7.mlp.c_fc.weight\n",
      "module.transformer.resblocks.7.mlp.c_fc.bias\n",
      "module.transformer.resblocks.7.mlp.c_proj.weight\n",
      "module.transformer.resblocks.7.mlp.c_proj.bias\n",
      "module.transformer.resblocks.8.ln_1.weight\n",
      "module.transformer.resblocks.8.ln_1.bias\n",
      "module.transformer.resblocks.8.attn.in_proj_weight\n",
      "module.transformer.resblocks.8.attn.in_proj_bias\n",
      "module.transformer.resblocks.8.attn.out_proj.weight\n",
      "module.transformer.resblocks.8.attn.out_proj.bias\n",
      "module.transformer.resblocks.8.ln_2.weight\n",
      "module.transformer.resblocks.8.ln_2.bias\n",
      "module.transformer.resblocks.8.mlp.c_fc.weight\n",
      "module.transformer.resblocks.8.mlp.c_fc.bias\n",
      "module.transformer.resblocks.8.mlp.c_proj.weight\n",
      "module.transformer.resblocks.8.mlp.c_proj.bias\n",
      "module.transformer.resblocks.9.ln_1.weight\n",
      "module.transformer.resblocks.9.ln_1.bias\n",
      "module.transformer.resblocks.9.attn.in_proj_weight\n",
      "module.transformer.resblocks.9.attn.in_proj_bias\n",
      "module.transformer.resblocks.9.attn.out_proj.weight\n",
      "module.transformer.resblocks.9.attn.out_proj.bias\n",
      "module.transformer.resblocks.9.ln_2.weight\n",
      "module.transformer.resblocks.9.ln_2.bias\n",
      "module.transformer.resblocks.9.mlp.c_fc.weight\n",
      "module.transformer.resblocks.9.mlp.c_fc.bias\n",
      "module.transformer.resblocks.9.mlp.c_proj.weight\n",
      "module.transformer.resblocks.9.mlp.c_proj.bias\n",
      "module.transformer.resblocks.10.ln_1.weight\n",
      "module.transformer.resblocks.10.ln_1.bias\n",
      "module.transformer.resblocks.10.attn.in_proj_weight\n",
      "module.transformer.resblocks.10.attn.in_proj_bias\n",
      "module.transformer.resblocks.10.attn.out_proj.weight\n",
      "module.transformer.resblocks.10.attn.out_proj.bias\n",
      "module.transformer.resblocks.10.ln_2.weight\n",
      "module.transformer.resblocks.10.ln_2.bias\n",
      "module.transformer.resblocks.10.mlp.c_fc.weight\n",
      "module.transformer.resblocks.10.mlp.c_fc.bias\n",
      "module.transformer.resblocks.10.mlp.c_proj.weight\n",
      "module.transformer.resblocks.10.mlp.c_proj.bias\n",
      "module.transformer.resblocks.11.ln_1.weight\n",
      "module.transformer.resblocks.11.ln_1.bias\n",
      "module.transformer.resblocks.11.attn.in_proj_weight\n",
      "module.transformer.resblocks.11.attn.in_proj_bias\n",
      "module.transformer.resblocks.11.attn.out_proj.weight\n",
      "module.transformer.resblocks.11.attn.out_proj.bias\n",
      "module.transformer.resblocks.11.ln_2.weight\n",
      "module.transformer.resblocks.11.ln_2.bias\n",
      "module.transformer.resblocks.11.mlp.c_fc.weight\n",
      "module.transformer.resblocks.11.mlp.c_fc.bias\n",
      "module.transformer.resblocks.11.mlp.c_proj.weight\n",
      "module.transformer.resblocks.11.mlp.c_proj.bias\n",
      "module.token_embedding.weight\n",
      "module.ln_final.weight\n",
      "module.ln_final.bias\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "depth = 24\n",
    "ckpt_path = \"/mnt/sda1/tkerssies/beit3_large_patch16_224.pth\"\n",
    "\n",
    "ckpt = torch.load(ckpt_path)[\"model\"]\n",
    "\n",
    "ckpt = {k.replace(\"beit3.\", \"\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k: v for k, v in ckpt.items() if not k.startswith(\"text_embed\")}\n",
    "ckpt = {k.replace(\"vision_embed.\", \"patch_embed.\"): v for k, v in ckpt.items()}\n",
    "ckpt.pop(\"patch_embed.mask_token\")\n",
    "ckpt.pop(\"mlm_head.weight\")\n",
    "ckpt.pop(\"mlm_head.bias\")\n",
    "ckpt.pop(\"mim_head.weight\")\n",
    "ckpt.pop(\"mim_head.bias\")\n",
    "ckpt = {k.replace(\"patch_embed.cls_token\", \"cls_token\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"encoder.\", \"\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"layers\", \"blocks\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k: v for k, v in ckpt.items() if \".B.\" not in k}\n",
    "ckpt = {k.replace(\".A.\", \".\"): v for k, v in ckpt.items()}\n",
    "ckpt[\"pos_embed\"] = ckpt.pop(\"embed_positions.weight\")\n",
    "ckpt = {k.replace(\"self_attn_layer_norm\", \"norm1\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"self_attn\", \"attn\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"inner_attn_ln\", \"proj.0\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"out_proj\", \"proj.1\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"ffn\", \"mlp\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"final_layer_norm\", \"norm2\"): v for k, v in ckpt.items()}\n",
    "ckpt = {k.replace(\"mlp_layernorm\", \"norm\"): v for k, v in ckpt.items()}\n",
    "ckpt[\"norm.weight\"] = ckpt.pop(\"layer_norm.weight\")\n",
    "ckpt[\"norm.bias\"] = ckpt.pop(\"layer_norm.bias\")\n",
    "ckpt[\"pos_embed\"] = ckpt[\"pos_embed\"][None, 2:, :]\n",
    "\n",
    "for block in range(depth):\n",
    "    q_key = f\"blocks.{block}.attn.q_proj\"\n",
    "    k_key = f\"blocks.{block}.attn.k_proj\"\n",
    "    v_key = f\"blocks.{block}.attn.v_proj\"\n",
    "\n",
    "    q_w = ckpt.pop(f\"{q_key}.weight\")\n",
    "    k_w = ckpt.pop(f\"{k_key}.weight\")\n",
    "    v_w = ckpt.pop(f\"{v_key}.weight\")\n",
    "\n",
    "    q_b = ckpt.pop(f\"{q_key}.bias\")\n",
    "    k_b = ckpt.pop(f\"{k_key}.bias\")\n",
    "    v_b = ckpt.pop(f\"{v_key}.bias\")\n",
    "\n",
    "    qkv_w = torch.cat([q_w, k_w, v_w], dim=0)\n",
    "    qkv_b = torch.cat([q_b, k_b, v_b])\n",
    "\n",
    "    ckpt[f\"blocks.{block}.attn.qkv.weight\"] = qkv_w\n",
    "    ckpt[f\"blocks.{block}.attn.qkv.bias\"] = qkv_b\n",
    "\n",
    "print(\"\\n\".join(ckpt.keys()))\n",
    "\n",
    "torch.save(ckpt, ckpt_path + \".timm\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "foundation-generalization",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}