{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install timm","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import timm","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchvision import models\nfrom torchvision import transforms\nimport torch.nn.functional as F","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MyModel(nn.Module):\n    def __init__(self, model_name, target_size=[224, 224]):\n        super().__init__()\n        self.target_size = target_size\n\n        self.encoder = timm.create_model(model_name, pretrained=True, num_classes=0)    \n        self.final = nn.AdaptiveAvgPool1d(64)\n        \n    def forward(self, x):\n        x = transforms.functional.resize(x, size=self.target_size, interpolation=transforms.InterpolationMode.BICUBIC)\n        x = x/255.0\n        x = transforms.functional.normalize(x,\n                                            mean=[0.48145466, 0.4578275, 0.40821073], \n                                            std=[0.26862954, 0.26130258, 0.27577711])\n        \n        x = self.encoder(x)\n        x = self.final(x)\n#         if self.normalize:\n#             x = torch.nn.functional.normalize(x)\n\n        return x","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = MyModel(\"vit_base_patch32_224_in21k\", target_size=[224, 224])\nmodel.eval()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# x = torch.randn(1, 3, 4545, 2324)\n# print(x.shape)\n# # Let's print it\n# model(x).shape","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"saved_model = torch.jit.script(model)\nsaved_model.save('saved_model.pt')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from zipfile import ZipFile\n\nwith ZipFile('submission.zip','w') as zip:\n    zip.write('saved_model.pt', arcname='saved_model.pt')","metadata":{},"execution_count":null,"outputs":[]}]}