{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":6799,"databundleVersionId":4225553,"sourceType":"competition"},{"sourceId":7790426,"sourceType":"datasetVersion","datasetId":4559896}],"dockerImageVersionId":30665,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\nfrom skimage import io\nimport numpy as np\nimport albumentations as A\nimport matplotlib.pyplot as plt","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:34:34.489017Z","iopub.execute_input":"2024-03-11T16:34:34.489398Z","iopub.status.idle":"2024-03-11T16:34:38.931947Z","shell.execute_reply.started":"2024-03-11T16:34:34.489368Z","shell.execute_reply":"2024-03-11T16:34:38.931009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_path = \"/kaggle/input/image-net-visualization/bird_1.png\"\nmodel_path = \"vit_small_patch16_224.dino\"\nimage_size = 1024","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:19.952014Z","iopub.execute_input":"2024-03-11T16:37:19.952419Z","iopub.status.idle":"2024-03-11T16:37:19.957159Z","shell.execute_reply.started":"2024-03-11T16:37:19.952389Z","shell.execute_reply":"2024-03-11T16:37:19.956147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_image(im_path, shape=image_size):\n    \n    \n    img = io.imread(im_path)[:,:,:3] ## HXWXC\n    H,W,C = img.shape\n    org_img = img.copy()\n    \n    img = A.Normalize()(image=img)[\"image\"] ## HxWxC\n    norm_img = img.copy()\n    \n    img = np.transpose(img,(2,0,1)) ## CxHxW\n    img = np.expand_dims(img,0) ## B=1,CxHxW\n    \n    img = torch.tensor(img)\n    img = nn.Upsample((shape,shape),mode='bilinear')(img)\n    \n    return img, norm_img, org_img","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:35:19.900909Z","iopub.execute_input":"2024-03-11T16:35:19.901645Z","iopub.status.idle":"2024-03-11T16:35:19.90918Z","shell.execute_reply.started":"2024-03-11T16:35:19.901614Z","shell.execute_reply":"2024-03-11T16:35:19.90821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img,norm_img,org_img = get_image(img_path)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:35:30.816121Z","iopub.execute_input":"2024-03-11T16:35:30.816866Z","iopub.status.idle":"2024-03-11T16:35:30.854659Z","shell.execute_reply.started":"2024-03-11T16:35:30.816834Z","shell.execute_reply":"2024-03-11T16:35:30.853728Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"io.imshow(org_img)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:35:31.053555Z","iopub.execute_input":"2024-03-11T16:35:31.053901Z","iopub.status.idle":"2024-03-11T16:35:31.413342Z","shell.execute_reply.started":"2024-03-11T16:35:31.053876Z","shell.execute_reply":"2024-03-11T16:35:31.412419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"io.imshow(norm_img)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:35:39.381807Z","iopub.execute_input":"2024-03-11T16:35:39.382527Z","iopub.status.idle":"2024-03-11T16:35:39.737951Z","shell.execute_reply.started":"2024-03-11T16:35:39.382498Z","shell.execute_reply":"2024-03-11T16:35:39.737015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n! pip install timm","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:35:42.745551Z","iopub.execute_input":"2024-03-11T16:35:42.745951Z","iopub.status.idle":"2024-03-11T16:35:55.760011Z","shell.execute_reply.started":"2024-03-11T16:35:42.745923Z","shell.execute_reply":"2024-03-11T16:35:55.758612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import timm\nmodel = timm.create_model(model_path,pretrained=True,\n                         img_size=image_size,\n                         dynamic_img_size=True)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:27.798406Z","iopub.execute_input":"2024-03-11T16:37:27.799208Z","iopub.status.idle":"2024-03-11T16:37:28.874518Z","shell.execute_reply.started":"2024-03-11T16:37:27.799175Z","shell.execute_reply":"2024-03-11T16:37:28.873583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = model.cuda()","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:29.730323Z","iopub.execute_input":"2024-03-11T16:37:29.731075Z","iopub.status.idle":"2024-03-11T16:37:29.767881Z","shell.execute_reply.started":"2024-03-11T16:37:29.731042Z","shell.execute_reply":"2024-03-11T16:37:29.76712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs = {}\ndef get_outputs(name:str):\n    def hook(model, input, output):\n        outputs[name] = output.detach()\n        \n    return hook","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:33.279702Z","iopub.execute_input":"2024-03-11T16:37:33.280041Z","iopub.status.idle":"2024-03-11T16:37:33.285225Z","shell.execute_reply.started":"2024-03-11T16:37:33.280018Z","shell.execute_reply":"2024-03-11T16:37:33.284214Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.blocks[-1].attn.q_norm.register_forward_hook(get_outputs('Q'))\nmodel.blocks[-1].attn.k_norm.register_forward_hook(get_outputs('K'))\nmodel.blocks[-1].register_forward_hook(get_outputs('features'))\n\nmodel(img.cuda())","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:34.496781Z","iopub.execute_input":"2024-03-11T16:37:34.497149Z","iopub.status.idle":"2024-03-11T16:37:34.628951Z","shell.execute_reply.started":"2024-03-11T16:37:34.49712Z","shell.execute_reply":"2024-03-11T16:37:34.628028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"scale = model.blocks[-1].attn.scale","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:35.972695Z","iopub.execute_input":"2024-03-11T16:37:35.973094Z","iopub.status.idle":"2024-03-11T16:37:35.977544Z","shell.execute_reply.started":"2024-03-11T16:37:35.973066Z","shell.execute_reply":"2024-03-11T16:37:35.976658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs[\"attention\"] = (outputs[\"Q\"]@ outputs[\"K\"].transpose(-2, -1))#*scale\n# outputs[\"attention\"] = outputs[\"attention\"].softmax(dim=-1)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:38.512485Z","iopub.execute_input":"2024-03-11T16:37:38.512877Z","iopub.status.idle":"2024-03-11T16:37:38.518051Z","shell.execute_reply.started":"2024-03-11T16:37:38.51285Z","shell.execute_reply":"2024-03-11T16:37:38.517097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs[\"attention\"].shape","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:52.13326Z","iopub.execute_input":"2024-03-11T16:37:52.134249Z","iopub.status.idle":"2024-03-11T16:37:52.140005Z","shell.execute_reply.started":"2024-03-11T16:37:52.134217Z","shell.execute_reply":"2024-03-11T16:37:52.139148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"b,num_heads,num_patches_1,num_patches_1 = outputs[\"attention\"].shape\nmap_size = int(np.sqrt(num_patches_1-1))\nfor attention_head in range(num_heads):\n    attention_map = outputs[\"attention\"][:,attention_head,0,1:] ## 1,4096\n    attention_map = attention_map.view(1,1,map_size,map_size)\n    \n    attention_map = nn.Upsample(size=(image_size,image_size))(attention_map)\n    \n    attention_map = attention_map[0,0,:,:].detach().cpu().numpy()\n    \n    io.imshow(attention_map)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:52.235613Z","iopub.execute_input":"2024-03-11T16:37:52.23589Z","iopub.status.idle":"2024-03-11T16:37:54.143441Z","shell.execute_reply.started":"2024-03-11T16:37:52.235868Z","shell.execute_reply":"2024-03-11T16:37:54.142544Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs[\"features\"].shape","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:37:54.145463Z","iopub.execute_input":"2024-03-11T16:37:54.14615Z","iopub.status.idle":"2024-03-11T16:37:54.152013Z","shell.execute_reply.started":"2024-03-11T16:37:54.146107Z","shell.execute_reply":"2024-03-11T16:37:54.151089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"features = outputs[\"features\"].mean(-1)\nfeatures = features[:,1:]\n\nfeatures = features.view(1,1,map_size,map_size)\n\nfeatures = nn.Upsample(size=(image_size,image_size))(features)\n\nfeatures = features[0,0,:,:].detach().cpu().numpy()\n\nio.imshow(features)","metadata":{"execution":{"iopub.status.busy":"2024-03-11T16:38:08.281915Z","iopub.execute_input":"2024-03-11T16:38:08.282838Z"},"trusted":true},"execution_count":null,"outputs":[]}]}