{"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":"!pip install timm\n!pip install fairseq","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:33:47.611676Z","iopub.execute_input":"2021-06-03T11:33:47.612049Z","iopub.status.idle":"2021-06-03T11:34:05.771306Z","shell.execute_reply.started":"2021-06-03T11:33:47.611971Z","shell.execute_reply":"2021-06-03T11:34:05.770404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\" Swin Transformer\nA PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`\n - https://arxiv.org/pdf/2103.14030\nCode/weights from https://github.com/microsoft/Swin-Transformer, original copyright/license info below\n\"\"\"\n# --------------------------------------------------------\n# Swin Transformer\n# Copyright (c) 2021 Microsoft\n# Licensed under The MIT License [see LICENSE for details]\n# Written by Ze Liu\n# --------------------------------------------------------\nimport logging\nimport math\nfrom copy import deepcopy\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint as checkpoint\n\nfrom timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD\nfrom timm.models.helpers import build_model_with_cfg, overlay_external_default_cfg\nfrom timm.models.layers import DropPath, to_2tuple, trunc_normal_\nfrom timm.models.registry import register_model\nfrom timm.models.vision_transformer import checkpoint_filter_fn, Mlp, _init_vit_weights\n\n_logger = logging.getLogger(__name__)\n\n\ndef _cfg(url='', **kwargs):\n return {\n 'url': url,\n 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,\n 'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,\n 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,\n 'first_conv': 'patch_embed.proj', 'classifier': 'head',\n **kwargs\n }\n\n\ndefault_cfgs_swin = {\n # patch models (my experiments)\n 'swin_base_patch4_window12_384': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22kto1k.pth',\n input_size=(3, 384, 384), crop_pct=1.0),\n\n 'swin_base_patch4_window7_224': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22kto1k.pth',\n ),\n\n 'swin_large_patch4_window12_384': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22kto1k.pth',\n input_size=(3, 384, 384), crop_pct=1.0),\n\n 'swin_large_patch4_window7_224': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22kto1k.pth',\n ),\n\n 'swin_small_patch4_window7_224': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth',\n ),\n\n 'swin_tiny_patch4_window7_224': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth',\n ),\n\n 'swin_base_patch4_window12_384_in22k': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth',\n input_size=(3, 384, 384), crop_pct=1.0, num_classes=21841),\n\n 'swin_base_patch4_window7_224_in22k': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth',\n num_classes=21841),\n\n 'swin_large_patch4_window12_384_in22k': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth',\n input_size=(3, 384, 384), crop_pct=1.0, num_classes=21841),\n\n 'swin_large_patch4_window7_224_in22k': _cfg(\n url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth',\n num_classes=21841),\n\n}\n\n\ndef window_partition(x, window_size: int):\n \"\"\"\n Args:\n x: (B, H, W, C)\n window_size (int): window size\n Returns:\n windows: (num_windows*B, window_size, window_size, C)\n \"\"\"\n B, H, W, C = x.shape\n x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)\n return windows\n\n\ndef window_reverse(windows, window_size: int, H: int, W: int):\n \"\"\"\n Args:\n windows: (num_windows*B, window_size, window_size, C)\n window_size (int): Window size\n H (int): Height of image\n W (int): Width of image\n Returns:\n x: (B, H, W, C)\n \"\"\"\n B = int(windows.shape[0] / (H * W / window_size / window_size))\n x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)\n return x\n\n\nclass WindowAttention(nn.Module):\n r\"\"\" Window based multi-head self attention (W-MSA) module with relative position bias.\n It supports both of shifted and non-shifted window.\n Args:\n dim (int): Number of input channels.\n window_size (tuple[int]): The height and width of the window.\n num_heads (int): Number of attention heads.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set\n attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n \"\"\"\n\n def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):\n\n super().__init__()\n self.dim = dim\n self.window_size = window_size # Wh, Ww\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = qk_scale or head_dim ** -0.5\n\n # define a parameter table of relative position bias\n self.relative_position_bias_table = nn.Parameter(\n torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH\n\n # get pair-wise relative position index for each token inside the window\n coords_h = torch.arange(self.window_size[0])\n coords_w = torch.arange(self.window_size[1])\n coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww\n coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww\n relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww\n relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2\n relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0\n relative_coords[:, :, 1] += self.window_size[1] - 1\n relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww\n self.register_buffer(\"relative_position_index\", relative_position_index)\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim)\n self.proj_drop = nn.Dropout(proj_drop)\n\n trunc_normal_(self.relative_position_bias_table, std=.02)\n self.softmax = nn.Softmax(dim=-1)\n\n def forward(self, x, mask: Optional[torch.Tensor] = None):\n \"\"\"\n Args:\n x: input features with shape of (num_windows*B, N, C)\n mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n \"\"\"\n B_, N, C = x.shape\n qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)\n\n q = q * self.scale\n attn = (q @ k.transpose(-2, -1))\n\n relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(\n self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH\n relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww\n attn = attn + relative_position_bias.unsqueeze(0)\n\n if mask is not None:\n nW = mask.shape[0]\n attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n attn = attn.view(-1, self.num_heads, N, N)\n attn = self.softmax(attn)\n else:\n attn = self.softmax(attn)\n\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B_, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n \nclass PatchEmbed(nn.Module):\n \"\"\" Image to Patch Embedding\n \"\"\"\n def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=nn.LayerNorm):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n self.img_size = img_size\n self.patch_size = patch_size\n self.patch_grid = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])\n self.num_patches = self.patch_grid[0] * self.patch_grid[1]\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n def forward(self, x):\n B, C, H, W = x.shape\n # FIXME look at relaxing size constraints\n assert H == self.img_size[0] and W == self.img_size[1], \\\n f\"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]}).\"\n x = self.proj(x).flatten(2).transpose(1, 2)\n x = self.norm(x)\n return x\n \n\nclass SwinTransformerBlock(nn.Module):\n r\"\"\" Swin Transformer Block.\n Args:\n dim (int): Number of input channels.\n input_resolution (tuple[int]): Input resulotion.\n num_heads (int): Number of attention heads.\n window_size (int): Window size.\n shift_size (int): Shift size for SW-MSA.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n drop (float, optional): Dropout rate. Default: 0.0\n attn_drop (float, optional): Attention dropout rate. Default: 0.0\n drop_path (float, optional): Stochastic depth rate. Default: 0.0\n act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n \"\"\"\n\n def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,\n mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,\n act_layer=nn.GELU, norm_layer=nn.LayerNorm):\n super().__init__()\n self.dim = dim\n self.input_resolution = input_resolution\n self.num_heads = num_heads\n self.window_size = window_size\n self.shift_size = shift_size\n self.mlp_ratio = mlp_ratio\n if min(self.input_resolution) <= self.window_size:\n # if window size is larger than input resolution, we don't partition windows\n self.shift_size = 0\n self.window_size = min(self.input_resolution)\n assert 0 <= self.shift_size < self.window_size, \"shift_size must in 0-window_size\"\n\n self.norm1 = norm_layer(dim)\n self.attn = WindowAttention(\n dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,\n qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)\n\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n if self.shift_size > 0:\n # calculate attention mask for SW-MSA\n H, W = self.input_resolution\n img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1\n h_slices = (slice(0, -self.window_size),\n slice(-self.window_size, -self.shift_size),\n slice(-self.shift_size, None))\n w_slices = (slice(0, -self.window_size),\n slice(-self.window_size, -self.shift_size),\n slice(-self.shift_size, None))\n cnt = 0\n for h in h_slices:\n for w in w_slices:\n img_mask[:, h, w, :] = cnt\n cnt += 1\n\n mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1\n mask_windows = mask_windows.view(-1, self.window_size * self.window_size)\n attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n else:\n attn_mask = None\n\n self.register_buffer(\"attn_mask\", attn_mask)\n\n def forward(self, x):\n H, W = self.input_resolution\n B, L, C = x.shape\n assert L == H * W, \"input feature has wrong size\"\n\n shortcut = x\n x = self.norm1(x)\n x = x.view(B, H, W, C)\n\n # cyclic shift\n if self.shift_size > 0:\n shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))\n else:\n shifted_x = x\n\n # partition windows\n x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C\n x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C\n\n # W-MSA/SW-MSA\n attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C\n\n # merge windows\n attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)\n shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C\n\n # reverse cyclic shift\n if self.shift_size > 0:\n x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))\n else:\n x = shifted_x\n x = x.view(B, H * W, C)\n\n # FFN\n x = shortcut + self.drop_path(x)\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n\n return x\n\n\nclass PatchMerging(nn.Module):\n r\"\"\" Patch Merging Layer.\n Args:\n input_resolution (tuple[int]): Resolution of input feature.\n dim (int): Number of input channels.\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n \"\"\"\n\n def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):\n super().__init__()\n self.input_resolution = input_resolution\n self.dim = dim\n self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n self.norm = norm_layer(4 * dim)\n\n def forward(self, x):\n \"\"\"\n x: B, H*W, C\n \"\"\"\n H, W = self.input_resolution\n B, L, C = x.shape\n assert L == H * W, \"input feature has wrong size\"\n assert H % 2 == 0 and W % 2 == 0, f\"x size ({H}*{W}) are not even.\"\n\n x = x.view(B, H, W, C)\n\n x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C\n x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C\n x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C\n x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C\n x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C\n x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C\n\n x = self.norm(x)\n x = self.reduction(x)\n\n return x\n\n def extra_repr(self) -> str:\n return f\"input_resolution={self.input_resolution}, dim={self.dim}\"\n\n def flops(self):\n H, W = self.input_resolution\n flops = H * W * self.dim\n flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim\n return flops\n\n\nclass BasicLayer(nn.Module):\n \"\"\" A basic Swin Transformer layer for one stage.\n Args:\n dim (int): Number of input channels.\n input_resolution (tuple[int]): Input resolution.\n depth (int): Number of blocks.\n num_heads (int): Number of attention heads.\n window_size (int): Local window size.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n drop (float, optional): Dropout rate. Default: 0.0\n attn_drop (float, optional): Attention dropout rate. Default: 0.0\n drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n \"\"\"\n\n def __init__(self, dim, input_resolution, depth, num_heads, window_size,\n mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,\n drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):\n\n super().__init__()\n self.dim = dim\n self.input_resolution = input_resolution\n self.depth = depth\n self.use_checkpoint = use_checkpoint\n\n # build blocks\n self.blocks = nn.ModuleList([\n SwinTransformerBlock(dim=dim, \n input_resolution=input_resolution,\n num_heads=num_heads, \n window_size=window_size,\n shift_size=0 if (i % 2 == 0) else window_size // 2,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias, \n qk_scale=qk_scale,\n drop=drop, \n attn_drop=attn_drop,\n drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n norm_layer=norm_layer\n )\n for i in range(depth)])\n\n # patch merging layer\n if downsample is not None:\n self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)\n else:\n self.downsample = None\n\n def forward(self, x):\n for blk in self.blocks:\n if not torch.jit.is_scripting() and self.use_checkpoint:\n x = checkpoint.checkpoint(blk, x)\n else:\n x = blk(x)\n if self.downsample is not None:\n x = self.downsample(x)\n return x\n\n def extra_repr(self) -> str:\n return f\"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}\"\n\n\nclass SwinTransformer(nn.Module):\n r\"\"\" Swin Transformer\n A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` -\n https://arxiv.org/pdf/2103.14030\n Args:\n img_size (int | tuple(int)): Input image size. Default 224\n patch_size (int | tuple(int)): Patch size. Default: 4\n in_chans (int): Number of input image channels. Default: 3\n num_classes (int): Number of classes for classification head. Default: 1000\n embed_dim (int): Patch embedding dimension. Default: 96\n depths (tuple(int)): Depth of each Swin Transformer layer.\n num_heads (tuple(int)): Number of attention heads in different layers.\n window_size (int): Window size. Default: 7\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None\n drop_rate (float): Dropout rate. Default: 0\n attn_drop_rate (float): Attention dropout rate. Default: 0\n drop_path_rate (float): Stochastic depth rate. Default: 0.1\n norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n ape (bool): If True, add absolute position embedding to the patch embedding. Default: False\n patch_norm (bool): If True, add normalization after patch embedding. Default: True\n use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n \"\"\"\n\n def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,\n embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24),\n window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,\n drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n norm_layer=nn.LayerNorm, ape=False, patch_norm=True,\n use_checkpoint=False, weight_init='', **kwargs):\n super().__init__()\n\n self.num_classes = num_classes\n self.num_layers = len(depths)\n self.embed_dim = embed_dim\n self.ape = ape\n self.patch_norm = patch_norm\n self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))\n self.mlp_ratio = mlp_ratio\n\n # split image into non-overlapping patches\n self.patch_embed = PatchEmbed(\n img_size=img_size, \n patch_size=patch_size, \n in_chans=in_chans, \n embed_dim=embed_dim,\n norm_layer=norm_layer if self.patch_norm else None\n )\n num_patches = self.patch_embed.num_patches\n self.patch_grid = self.patch_embed.patch_grid\n\n # absolute position embedding\n if self.ape:\n self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))\n trunc_normal_(self.absolute_pos_embed, std=.02)\n else:\n self.absolute_pos_embed = None\n\n self.pos_drop = nn.Dropout(p=drop_rate)\n\n # stochastic depth\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule\n\n # build layers\n layers = []\n for i_layer in range(self.num_layers):\n layers += [BasicLayer(\n dim=int(embed_dim * 2 ** i_layer),\n input_resolution=(self.patch_grid[0] // (2 ** i_layer), self.patch_grid[1] // (2 ** i_layer)),\n depth=depths[i_layer],\n num_heads=num_heads[i_layer],\n window_size=window_size,\n mlp_ratio=self.mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate,\n drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],\n norm_layer=norm_layer,\n downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,\n use_checkpoint=use_checkpoint)\n ]\n self.layers = nn.Sequential(*layers)\n\n self.norm = norm_layer(self.num_features)\n self.avgpool = nn.AdaptiveAvgPool1d(1)\n self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()\n\n assert weight_init in ('jax', 'jax_nlhb', 'nlhb', '')\n head_bias = -math.log(self.num_classes) if 'nlhb' in weight_init else 0.\n if weight_init.startswith('jax'):\n for n, m in self.named_modules():\n _init_vit_weights(m, n, head_bias=head_bias, jax_impl=True)\n else:\n self.apply(_init_vit_weights)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'absolute_pos_embed'}\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {'relative_position_bias_table'}\n\n def forward_features(self, x):\n x = self.patch_embed(x)\n if self.absolute_pos_embed is not None:\n x = x + self.absolute_pos_embed\n x = self.pos_drop(x)\n x = self.layers(x)\n x = self.norm(x) # B L C\n #x = self.avgpool(x.transpose(1, 2)) # B C 1\n #x = torch.flatten(x, 1)\n return x\n\n def forward(self, x):\n x = self.forward_features(x)\n x = self.head(x)\n return x\n\n\ndef _create_swin_transformer(variant, pretrained=False, default_cfg=None, **kwargs):\n if default_cfg is None:\n default_cfg = deepcopy(default_cfgs_swin[variant])\n overlay_external_default_cfg(default_cfg, kwargs)\n default_num_classes = default_cfg['num_classes']\n default_img_size = default_cfg['input_size'][-2:]\n\n num_classes = kwargs.pop('num_classes', default_num_classes)\n img_size = kwargs.pop('img_size', default_img_size)\n if kwargs.get('features_only', None):\n raise RuntimeError('features_only not implemented for Vision Transformer models.')\n\n model = build_model_with_cfg(\n SwinTransformer, \n variant, pretrained,\n default_cfg=default_cfg,\n img_size=img_size,\n num_classes=num_classes,\n pretrained_filter_fn=checkpoint_filter_fn,\n **kwargs)\n\n return model\n\n\n\n@register_model\ndef swin_base_patch4_window12_384(pretrained=False, **kwargs):\n \"\"\" Swin-B @ 384x384, pretrained ImageNet-22k, fine tune 1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)\n return _create_swin_transformer('swin_base_patch4_window12_384', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_base_patch4_window7_224(pretrained=False, **kwargs):\n \"\"\" Swin-B @ 224x224, pretrained ImageNet-22k, fine tune 1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)\n return _create_swin_transformer('swin_base_patch4_window7_224', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_large_patch4_window12_384(pretrained=False, **kwargs):\n \"\"\" Swin-L @ 384x384, pretrained ImageNet-22k, fine tune 1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)\n return _create_swin_transformer('swin_large_patch4_window12_384', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_large_patch4_window7_224(pretrained=False, **kwargs):\n \"\"\" Swin-L @ 224x224, pretrained ImageNet-22k, fine tune 1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)\n return _create_swin_transformer('swin_large_patch4_window7_224', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_small_patch4_window7_224(pretrained=False, **kwargs):\n \"\"\" Swin-S @ 224x224, trained ImageNet-1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 18, 2), num_heads=(3, 6, 12, 24), **kwargs)\n return _create_swin_transformer('swin_small_patch4_window7_224', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_tiny_patch4_window7_224(pretrained=False, **kwargs):\n \"\"\" Swin-T @ 224x224, trained ImageNet-1k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24), **kwargs)\n return _create_swin_transformer('swin_tiny_patch4_window7_224', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_base_patch4_window12_384_in22k(pretrained=False, **kwargs):\n \"\"\" Swin-B @ 384x384, trained ImageNet-22k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)\n return _create_swin_transformer('swin_base_patch4_window12_384_in22k', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_base_patch4_window7_224_in22k(pretrained=False, **kwargs):\n \"\"\" Swin-B @ 224x224, trained ImageNet-22k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)\n return _create_swin_transformer('swin_base_patch4_window7_224_in22k', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_large_patch4_window12_384_in22k(pretrained=False, **kwargs):\n \"\"\" Swin-L @ 384x384, trained ImageNet-22k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)\n return _create_swin_transformer('swin_large_patch4_window12_384_in22k', pretrained=pretrained, **model_kwargs)\n\n\n@register_model\ndef swin_large_patch4_window7_224_in22k(pretrained=False, **kwargs):\n \"\"\" Swin-L @ 224x224, trained ImageNet-22k\n \"\"\"\n model_kwargs = dict(\n patch_size=4, window_size=7, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)\n return _create_swin_transformer('swin_large_patch4_window7_224_in22k', pretrained=pretrained, **model_kwargs)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:05.774987Z","iopub.execute_input":"2021-06-03T11:34:05.775260Z","iopub.status.idle":"2021-06-03T11:34:07.820042Z","shell.execute_reply.started":"2021-06-03T11:34:05.775231Z","shell.execute_reply":"2021-06-03T11:34:07.819115Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\" Vision Transformer (ViT) in PyTorch\nA PyTorch implement of Vision Transformers as described in\n'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929\nThe official jax code is released and available at https://github.com/google-research/vision_transformer\nDeiT model defs and weights from https://github.com/facebookresearch/deit,\npaper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877\nAcknowledgments:\n* The paper authors for releasing code and weights, thanks!\n* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out\nfor some einops/einsum fun\n* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT\n* Bert reference code checks against Huggingface Transformers and Tensorflow Bert\nHacked together by / Copyright 2020 Ross Wightman\n\"\"\"\nimport math\nimport logging\nfrom functools import partial\nfrom collections import OrderedDict\nfrom copy import deepcopy\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD\nfrom timm.models.helpers import load_pretrained, build_model_with_cfg, overlay_external_default_cfg\nfrom timm.models.layers import DropPath, to_2tuple, trunc_normal_, lecun_normal_\nfrom timm.models.vision_transformer import Mlp\nfrom timm.models.registry import register_model\n\n_logger = logging.getLogger(__name__)\n\n\ndef _cfg(url='', **kwargs):\n return {\n 'url': url,\n 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,\n 'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,\n 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,\n 'first_conv': 'patch_embed.proj', 'classifier': 'head',\n **kwargs\n }\n\n\ndefault_cfgs = {\n # patch models (my experiments)\n 'vit_small_patch16_224': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth',\n ),\n\n # patch models (weights ported from official Google JAX impl)\n 'vit_base_patch16_224': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',\n mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),\n ),\n 'vit_base_patch32_224': _cfg(\n url='', # no official model weights for this combo, only for in21k\n mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_base_patch16_384': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth',\n input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),\n 'vit_base_patch32_384': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p32_384-830016f5.pth',\n input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),\n 'vit_large_patch16_224': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',\n mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_large_patch32_224': _cfg(\n url='', # no official model weights for this combo, only for in21k\n mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_large_patch16_384': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_384-b3be5167.pth',\n input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),\n 'vit_large_patch32_384': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pth',\n input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0),\n\n # patch models, imagenet21k (weights ported from official Google JAX impl)\n 'vit_base_patch16_224_in21k': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth',\n num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_base_patch32_224_in21k': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth',\n num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_large_patch16_224_in21k': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth',\n num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_large_patch32_224_in21k': _cfg(\n url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth',\n num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n 'vit_huge_patch14_224_in21k': _cfg(\n hf_hub='timm/vit_huge_patch14_224_in21k',\n num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n\n # deit models (FB weights)\n 'vit_deit_tiny_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_tiny_patch16_224-a1311bcf.pth'),\n 'vit_deit_small_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth'),\n 'vit_deit_base_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth',),\n 'vit_deit_base_patch16_384': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_base_patch16_384-8de9b5d1.pth',\n input_size=(3, 384, 384), crop_pct=1.0),\n 'vit_deit_tiny_distilled_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_tiny_distilled_patch16_224-b40b3cf7.pth',\n classifier=('head', 'head_dist')),\n 'vit_deit_small_distilled_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_small_distilled_patch16_224-649709d9.pth',\n classifier=('head', 'head_dist')),\n 'vit_deit_base_distilled_patch16_224': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_base_distilled_patch16_224-df68dfff.pth',\n classifier=('head', 'head_dist')),\n 'vit_deit_base_distilled_patch16_384': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_base_distilled_patch16_384-d0272ac0.pth',\n input_size=(3, 384, 384), crop_pct=1.0, classifier=('head', 'head_dist')),\n 'vit_deit_base_distilled_patch16_448': _cfg(\n url='https://dl.fbaipublicfiles.com/deit/deit_base_distilled_patch16_384-d0272ac0.pth',\n input_size=(3, 448, 448), crop_pct=1.0, classifier=('head', 'head_dist')),\n \n # ViT ImageNet-21K-P pretraining\n 'vit_base_patch16_224_miil_in21k': _cfg(\n url='https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/timm/vit_base_patch16_224_in21k_miil.pth',\n mean=(0, 0, 0), std=(1, 1, 1), crop_pct=0.875, interpolation='bilinear', num_classes=11221,\n ),\n 'vit_base_patch16_224_miil': _cfg(\n url='https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/timm'\n '/vit_base_patch16_224_1k_miil_84_4.pth',\n mean=(0, 0, 0), std=(1, 1, 1), crop_pct=0.875, interpolation='bilinear',\n ),\n}\n\n\nclass Mlp(nn.Module):\n def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x):\n x = self.fc1(x)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n\n\nclass Attention(nn.Module):\n def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = qk_scale or head_dim ** -0.5\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x):\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)\n\n attn = (q @ k.transpose(-2, -1)) * self.scale\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n\nclass Block(nn.Module):\n\n def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):\n super().__init__()\n self.norm1 = norm_layer(dim)\n self.attn = Attention(\n dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)\n # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n def forward(self, x):\n x = x + self.drop_path(self.attn(self.norm1(x)))\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n return x\n\n\nclass PatchEmbed(nn.Module):\n \"\"\" Image to Patch Embedding\n \"\"\"\n def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n self.img_size = img_size\n self.patch_size = patch_size\n self.patch_grid = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])\n self.num_patches = self.patch_grid[0] * self.patch_grid[1]\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n def forward(self, x):\n B, C, H, W = x.shape\n # FIXME look at relaxing size constraints\n assert H == self.img_size[0] and W == self.img_size[1], \\\n f\"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]}).\"\n x = self.proj(x).flatten(2).transpose(1, 2)\n x = self.norm(x)\n return x\n\n\nclass VisionTransformer(nn.Module):\n \"\"\" Vision Transformer\n A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`\n - https://arxiv.org/abs/2010.11929\n Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`\n - https://arxiv.org/abs/2012.12877\n \"\"\"\n\n def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,\n num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None, distilled=False,\n drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,\n act_layer=None, weight_init=''):\n \"\"\"\n Args:\n img_size (int, tuple): input image size\n patch_size (int, tuple): patch size\n in_chans (int): number of input channels\n num_classes (int): number of classes for classification head\n embed_dim (int): embedding dimension\n depth (int): depth of transformer\n num_heads (int): number of attention heads\n mlp_ratio (int): ratio of mlp hidden dim to embedding dim\n qkv_bias (bool): enable bias for qkv if True\n qk_scale (float): override default qk scale of head_dim ** -0.5 if set\n representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set\n distilled (bool): model includes a distillation token and head as in DeiT models\n drop_rate (float): dropout rate\n attn_drop_rate (float): attention dropout rate\n drop_path_rate (float): stochastic depth rate\n embed_layer (nn.Module): patch embedding layer\n norm_layer: (nn.Module): normalization layer\n weight_init: (str): weight init scheme\n \"\"\"\n super().__init__()\n self.num_classes = num_classes\n self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models\n self.num_tokens = 2 if distilled else 1\n norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)\n act_layer = act_layer or nn.GELU\n\n self.patch_embed = embed_layer(\n img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)\n num_patches = self.patch_embed.num_patches\n\n self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None\n self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))\n self.pos_drop = nn.Dropout(p=drop_rate)\n\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule\n self.blocks = nn.Sequential(*[\n Block(\n dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer)\n for i in range(depth)])\n self.norm = norm_layer(embed_dim)\n\n # Representation layer\n if representation_size and not distilled:\n self.num_features = representation_size\n self.pre_logits = nn.Sequential(OrderedDict([\n ('fc', nn.Linear(embed_dim, representation_size)),\n ('act', nn.Tanh())\n ]))\n else:\n self.pre_logits = nn.Identity()\n\n # Classifier head(s)\n self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()\n self.head_dist = None\n if distilled:\n self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()\n\n # Weight init\n assert weight_init in ('jax', 'jax_nlhb', 'nlhb', '')\n head_bias = -math.log(self.num_classes) if 'nlhb' in weight_init else 0.\n trunc_normal_(self.pos_embed, std=.02)\n if self.dist_token is not None:\n trunc_normal_(self.dist_token, std=.02)\n if weight_init.startswith('jax'):\n # leave cls token as zeros to match jax impl\n for n, m in self.named_modules():\n _init_vit_weights(m, n, head_bias=head_bias, jax_impl=True)\n else:\n trunc_normal_(self.cls_token, std=.02)\n self.apply(_init_vit_weights)\n\n def _init_weights(self, m):\n # this fn left here for compat with downstream users\n _init_vit_weights(m)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'pos_embed', 'cls_token', 'dist_token'}\n\n def get_classifier(self):\n if self.dist_token is None:\n return self.head\n else:\n return self.head, self.head_dist\n\n def reset_classifier(self, num_classes, global_pool=''):\n self.num_classes = num_classes\n self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n if self.num_tokens == 2:\n self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()\n\n def forward_features(self, x):\n x = self.patch_embed(x)\n cls_token = self.cls_token.expand(x.shape[0], -1, -1) # stole cls_tokens impl from Phil Wang, thanks\n if self.dist_token is None:\n x = torch.cat((cls_token, x), dim=1)\n else:\n x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)\n x = self.pos_drop(x + self.pos_embed)\n #print(x.size()) torch.Size([8, 577, 1024])\n x = self.blocks(x)\n #print(x.size()) torch.Size([8, 577, 1024])\n x = self.norm(x)\n #print(x.size()) torch.Size([8, 577, 1024])\n return x\n #if self.dist_token is None:\n # return self.pre_logits(x[:, 0])\n #else:\n # return x[:, 0], x[:, 1]\n\n def forward(self, x):\n x = self.forward_features(x)\n if self.head_dist is not None:\n x, x_dist = self.head(x[0]), self.head_dist(x[1]) # x must be a tuple\n if self.training and not torch.jit.is_scripting():\n # during inference, return the average of both classifier predictions\n return x, x_dist\n else:\n return (x + x_dist) / 2\n else:\n x = self.head(x)\n return x\n\n\ndef _init_vit_weights(m, n: str = '', head_bias: float = 0., jax_impl: bool = False):\n \"\"\" ViT weight initialization\n * When called without n, head_bias, jax_impl args it will behave exactly the same\n as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).\n * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl\n \"\"\"\n if isinstance(m, nn.Linear):\n if n.startswith('head'):\n nn.init.zeros_(m.weight)\n nn.init.constant_(m.bias, head_bias)\n elif n.startswith('pre_logits'):\n lecun_normal_(m.weight)\n nn.init.zeros_(m.bias)\n else:\n if jax_impl:\n nn.init.xavier_uniform_(m.weight)\n if m.bias is not None:\n if 'mlp' in n:\n nn.init.normal_(m.bias, std=1e-6)\n else:\n nn.init.zeros_(m.bias)\n else:\n trunc_normal_(m.weight, std=.02)\n if m.bias is not None:\n nn.init.zeros_(m.bias)\n elif jax_impl and isinstance(m, nn.Conv2d):\n # NOTE conv was left to pytorch default in my original init\n lecun_normal_(m.weight)\n if m.bias is not None:\n nn.init.zeros_(m.bias)\n elif isinstance(m, nn.LayerNorm):\n nn.init.zeros_(m.bias)\n nn.init.ones_(m.weight)\n\n\ndef resize_pos_embed(posemb, posemb_new, num_tokens=1):\n # Rescale the grid of position embeddings when loading from state_dict. Adapted from\n # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224\n _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)\n ntok_new = posemb_new.shape[1]\n if num_tokens:\n posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:]\n ntok_new -= num_tokens\n else:\n posemb_tok, posemb_grid = posemb[:, :0], posemb[0]\n gs_old = int(math.sqrt(len(posemb_grid)))\n gs_new = int(math.sqrt(ntok_new))\n _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)\n posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)\n posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')\n posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)\n posemb = torch.cat([posemb_tok, posemb_grid], dim=1)\n return posemb\n\n\ndef checkpoint_filter_fn(state_dict, model):\n \"\"\" convert patch embedding weight from manual patchify + linear proj to conv\"\"\"\n out_dict = {}\n if 'model' in state_dict:\n # For deit models\n state_dict = state_dict['model']\n for k, v in state_dict.items():\n if 'patch_embed.proj.weight' in k and len(v.shape) < 4:\n # For old models that I trained prior to conv based patchification\n O, I, H, W = model.patch_embed.proj.weight.shape\n v = v.reshape(O, -1, H, W)\n elif k == 'pos_embed' and v.shape != model.pos_embed.shape:\n # To resize pos embedding when using model at different size from pretrained weights\n v = resize_pos_embed(v, model.pos_embed, getattr(model, 'num_tokens', 1))\n out_dict[k] = v\n return out_dict\n\n\ndef _create_vision_transformer(variant, pretrained=False, default_cfg=None, **kwargs):\n if default_cfg is None:\n default_cfg = deepcopy(default_cfgs[variant])\n overlay_external_default_cfg(default_cfg, kwargs)\n default_num_classes = default_cfg['num_classes']\n default_img_size = default_cfg['input_size'][-2:]\n\n num_classes = kwargs.pop('num_classes', default_num_classes)\n img_size = kwargs.pop('img_size', default_img_size)\n repr_size = kwargs.pop('representation_size', None)\n if repr_size is not None and num_classes != default_num_classes:\n # Remove representation layer if fine-tuning. This may not always be the desired action,\n # but I feel better than doing nothing by default for fine-tuning. Perhaps a better interface?\n _logger.warning(\"Removing representation layer for fine-tuning.\")\n repr_size = None\n\n if kwargs.get('features_only', None):\n raise RuntimeError('features_only not implemented for Vision Transformer models.')\n\n model = build_model_with_cfg(\n VisionTransformer, variant, pretrained,\n default_cfg=default_cfg,\n img_size=img_size,\n num_classes=num_classes,\n representation_size=repr_size,\n pretrained_filter_fn=checkpoint_filter_fn,\n **kwargs)\n\n return model\n\n\n@register_model\ndef vit_small_patch16_224(pretrained=False, **kwargs):\n \"\"\" My custom 'small' ViT model. embed_dim=768, depth=8, num_heads=8, mlp_ratio=3.\n NOTE:\n * this differs from the DeiT based 'small' definitions with embed_dim=384, depth=12, num_heads=6\n * this model does not have a bias for QKV (unlike the official ViT and DeiT models)\n \"\"\"\n model_kwargs = dict(\n patch_size=16, embed_dim=768, depth=8, num_heads=8, mlp_ratio=3.,\n qkv_bias=False, norm_layer=nn.LayerNorm, **kwargs)\n if pretrained:\n # NOTE my scale was wrong for original weights, leaving this here until I have better ones for this model\n model_kwargs.setdefault('qk_scale', 768 ** -0.5)\n model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch16_224(pretrained=False, **kwargs):\n \"\"\" ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_base_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch32_224(pretrained=False, **kwargs):\n \"\"\" ViT-Base (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929). No pretrained weights.\n \"\"\"\n model_kwargs = dict(patch_size=32, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_base_patch32_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch16_384(pretrained=False, **kwargs):\n \"\"\" ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_base_patch16_384', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch32_384(pretrained=False, **kwargs):\n \"\"\" ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=32, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_base_patch32_384', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch16_224(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)\n model = _create_vision_transformer('vit_large_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch32_224(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929). No pretrained weights.\n \"\"\"\n model_kwargs = dict(patch_size=32, embed_dim=1024, depth=24, num_heads=16, **kwargs)\n model = _create_vision_transformer('vit_large_patch32_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch16_384(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)\n model = _create_vision_transformer('vit_large_patch16_384', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch32_384(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(patch_size=32, embed_dim=1024, depth=24, num_heads=16, **kwargs)\n model = _create_vision_transformer('vit_large_patch32_384', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch16_224_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(\n patch_size=16, embed_dim=768, depth=12, num_heads=12, representation_size=768, **kwargs)\n model = _create_vision_transformer('vit_base_patch16_224_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch32_224_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(\n patch_size=32, embed_dim=768, depth=12, num_heads=12, representation_size=768, **kwargs)\n model = _create_vision_transformer('vit_base_patch32_224_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch16_224_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(\n patch_size=16, embed_dim=1024, depth=24, num_heads=16, representation_size=1024, **kwargs)\n model = _create_vision_transformer('vit_large_patch16_224_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_large_patch32_224_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.\n \"\"\"\n model_kwargs = dict(\n patch_size=32, embed_dim=1024, depth=24, num_heads=16, representation_size=1024, **kwargs)\n model = _create_vision_transformer('vit_large_patch32_224_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_huge_patch14_224_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).\n ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.\n NOTE: converted weights not currently available, too large for github release hosting.\n \"\"\"\n model_kwargs = dict(\n patch_size=14, embed_dim=1280, depth=32, num_heads=16, representation_size=1280, **kwargs)\n model = _create_vision_transformer('vit_huge_patch14_224_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_tiny_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT-tiny model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=192, depth=12, num_heads=3, **kwargs)\n model = _create_vision_transformer('vit_deit_tiny_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_small_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT-small model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)\n model = _create_vision_transformer('vit_deit_small_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_base_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT base model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_deit_base_patch16_224', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_base_patch16_384(pretrained=False, **kwargs):\n \"\"\" DeiT base model @ 384x384 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer('vit_deit_base_patch16_384', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_tiny_distilled_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT-tiny distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=192, depth=12, num_heads=3, **kwargs)\n model = _create_vision_transformer(\n 'vit_deit_tiny_distilled_patch16_224', pretrained=pretrained, distilled=True, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_small_distilled_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT-small distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)\n model = _create_vision_transformer(\n 'vit_deit_small_distilled_patch16_224', pretrained=pretrained, distilled=True, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_base_distilled_patch16_224(pretrained=False, **kwargs):\n \"\"\" DeiT-base distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer(\n 'vit_deit_base_distilled_patch16_224', pretrained=pretrained, distilled=True, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_base_distilled_patch16_384(pretrained=False, **kwargs):\n \"\"\" DeiT-base distilled model @ 384x384 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer(\n 'vit_deit_base_distilled_patch16_384', pretrained=pretrained, distilled=True, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_deit_base_distilled_patch16_448(pretrained=False, **kwargs):\n \"\"\" DeiT-base distilled model @ 384x384 from paper (https://arxiv.org/abs/2012.12877).\n ImageNet-1k weights from https://github.com/facebookresearch/deit.\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)\n model = _create_vision_transformer(\n 'vit_deit_base_distilled_patch16_448', pretrained=pretrained, distilled=True, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch16_224_miil_in21k(pretrained=False, **kwargs):\n \"\"\" ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).\n Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, qkv_bias=False, **kwargs)\n model = _create_vision_transformer('vit_base_patch16_224_miil_in21k', pretrained=pretrained, **model_kwargs)\n return model\n\n\n@register_model\ndef vit_base_patch16_224_miil(pretrained=False, **kwargs):\n \"\"\" ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).\n Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K\n \"\"\"\n model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, qkv_bias=False, **kwargs)\n model = _create_vision_transformer('vit_base_patch16_224_miil', pretrained=pretrained, **model_kwargs)\n return model","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:07.822405Z","iopub.execute_input":"2021-06-03T11:34:07.822746Z","iopub.status.idle":"2021-06-03T11:34:07.929032Z","shell.execute_reply.started":"2021-06-03T11:34:07.822706Z","shell.execute_reply":"2021-06-03T11:34:07.927853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pickle\nimport numpy as np\nimport Levenshtein\n\ndef read_pickle_from_file(pickle_file):\n with open(pickle_file,'rb') as f:\n x = pickle.load(f)\n return x\n\n\ndef write_pickle_to_file(pickle_file, x):\n with open(pickle_file, 'wb') as f:\n pickle.dump(x, f, pickle.HIGHEST_PROTOCOL)\n\n\nclass YNakamaTokenizer(object):\n\n def __init__(self, is_load=True):\n self.stoi = {}\n self.itos = {}\n\n if is_load:\n self.stoi = read_pickle_from_file('../input/bms-predictions/tokenizer.stoi.pickle')\n self.itos = {k: v for v, k in self.stoi.items()}\n\n def __len__(self):\n return len(self.stoi)\n\n def build_vocab(self, text):\n vocab = set()\n for t in text:\n vocab.update(t.split(' '))\n vocab = sorted(vocab)\n vocab.append('')\n vocab.append('')\n vocab.append('')\n for i, s in enumerate(vocab):\n self.stoi[s] = i\n self.itos = {k: v for v, k in self.stoi.items()}\n\n def one_text_to_sequence(self, text):\n sequence = []\n sequence.append(self.stoi[''])\n for s in text.split(' '):\n sequence.append(self.stoi[s])\n sequence.append(self.stoi[''])\n return sequence\n\n def one_sequence_to_text(self, sequence):\n return ''.join(list(map(lambda i: self.itos[i], sequence)))\n\n def one_predict_to_inchi(self, predict):\n inchi = 'InChI=1S/'\n for p in predict:\n if p == self.stoi[''] or p == self.stoi['']:\n break\n inchi += self.itos[p]\n return inchi\n\n def text_to_sequence(self, text):\n sequence = [\n self.one_text_to_sequence(t)\n for t in text\n ]\n return sequence\n\n def sequence_to_text(self, sequence):\n text = [\n self.one_sequence_to_text(s)\n for s in sequence\n ]\n return text\n\n def predict_to_inchi(self, predict):\n inchi = [\n self.one_predict_to_inchi(p)\n for p in predict\n ]\n return inchi\n\n\ndef compute_lb_score(predict, truth):\n score = []\n for p, t in zip(predict, truth):\n s = Levenshtein.distance(p, t)\n score.append(s)\n score = np.array(score)\n return score","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-03T11:34:07.930962Z","iopub.execute_input":"2021-06-03T11:34:07.931430Z","iopub.status.idle":"2021-06-03T11:34:07.954630Z","shell.execute_reply.started":"2021-06-03T11:34:07.931390Z","shell.execute_reply":"2021-06-03T11:34:07.953811Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from collections import defaultdict\nimport random\nimport pandas as pd\nimport cv2\n\nimport torch\nfrom torch.utils.data import Sampler\nfrom torch.utils.data.dataset import Dataset\n\n#from bms import *\nfrom albumentations import (\n Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, \n IAAAdditiveGaussianNoise, Transpose, Blur, RandomRotate90, RandomScale,\n )\n\n\nSTOI = {\n '': 190,\n '': 191,\n '': 192,\n}\n\nimage_size = 384\nvocab_size = 193\nmax_length = 280 #300 #275\n\n\ndata_dir = '../../tnt/raj/heng_clean_384/data/'\n\n\ndef read_pickle_from_file(pickle_file):\n with open(pickle_file, 'rb') as f:\n x = pickle.load(f)\n return x\n\n\ndef write_pickle_to_file(pickle_file, x):\n with open(pickle_file, 'wb') as f:\n pickle.dump(x, f, pickle.HIGHEST_PROTOCOL)\n\n\ndef pad_sequence_to_max_length(sequence, max_length, padding_value):\n batch_size =len(sequence)\n pad_sequence = np.full((batch_size, max_length), padding_value, np.int32)\n for b, s in enumerate(sequence):\n L = len(s)\n pad_sequence[b, :L, ...] = s\n return pad_sequence\n\n\ndef load_tokenizer():\n tokenizer = YNakamaTokenizer(is_load=True)\n print('len(tokenizer) : vocab_size', len(tokenizer))\n for k, v in STOI.items():\n assert tokenizer.stoi[k]==v\n return tokenizer\n\n\ndef make_fold(mode='train-1'):\n if 'train' in mode:\n df = read_pickle_from_file(data_dir+'/df_train.more.csv.pickle')\n #df = pd.read_csv(data_dir + 'train_labels.csv')\n df_fold = pd.read_csv(data_dir+'/df_fold.csv')\n df = df.merge(df_fold, on='image_id')\n df.loc[:, 'path'] = 'train'\n df.loc[:, 'orientation'] = 0\n\n df['fold'] = df['fold'].astype(int)\n\n fold = int(mode[-1])\n df_train = df[df.fold != fold].reset_index(drop=True)\n df_valid = df[df.fold == fold].reset_index(drop=True)\n return df_train, df_valid\n\n # Index(['image_id', 'InChI'], dtype='object')\n if 'test' in mode:\n #df = pd.read_csv(data_dir+'/sample_submission.csv')\n df = pd.read_csv('../input/bms-predictions/invalid_test_77_ichi.csv')\n df_orientation = pd.read_csv('../input/bms-predictions/test_orientation.csv')\n df = df.merge(df_orientation, on='image_id')\n\n df.loc[:, 'path'] = 'test'\n #df.loc[:, 'InChI'] = '0'\n df.loc[:, 'formula'] = '0'\n df.loc[:, 'text'] = '0'\n df.loc[:, 'sequence'] = pd.Series([[0]] * len(df))\n df.loc[:, 'length'] = df.InChI.str.len()\n\n df_test = df\n return df_test\n\n\n#####################################################################################################\nclass FixNumSampler(Sampler):\n def __init__(self, dataset, length=-1, is_shuffle=False):\n if length <= 0:\n length = len(dataset)\n\n self.is_shuffle = is_shuffle\n self.length = length\n\n def __iter__(self):\n index = np.arange(self.length)\n if self.is_shuffle: \n random.shuffle(index)\n return iter(index)\n\n def __len__(self):\n return self.length\n\n\n# see https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-inference/data\ndef remote_unrotate_augment(r):\n image = r['image']\n h, w = image.shape\n\n if h > w:\n image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)\n\n image = cv2.resize(image, dsize=(image_size,image_size), interpolation=cv2.INTER_LINEAR)\n assert image_size == 384\n\n r['image'] = image\n return r\n\n\ndef null_augment(r):\n image = r['image']\n image = cv2.resize(image, dsize=(image_size, image_size), interpolation=cv2.INTER_LINEAR)\n assert image_size==384\n r['image'] = image\n return r\n\n\ndef get_augmentation():\n transform = [\n #RandomRotate90(p=0.01),\n RandomScale(scale_limit=(-0.2, +0.2), interpolation=1, always_apply=False, p=0.3),\n Cutout(num_holes=100, max_h_size=1, max_w_size=1, always_apply=False, p=0.3),\n ]\n return Compose(transform)\n\n\ndef null_augment_tr(r):\n \n image = r['image']\n trans = get_augmentation()\n image = trans(image=image)['image']\n\n image = cv2.resize(image, dsize=(image_size,image_size), interpolation=cv2.INTER_LINEAR)\n assert image_size == 384\n \n r['image'] = image\n return r\n\n\nclass BmsDataset(Dataset):\n def __init__(self, df, tokenizer, augment=null_augment):\n super().__init__()\n self.tokenizer = tokenizer\n self.df = df\n self.augment = augment\n self.length = len(self.df)\n\n def __str__(self):\n string = ''\n string += '\\tlen = %d\\n'%len(self)\n string += '\\tdf = %s\\n'%str(self.df.shape)\n return string\n\n def __len__(self):\n return self.length\n\n def __getitem__(self, index):\n d = self.df.iloc[index]\n \n image_file = '../input/bms-molecular-translation' +'/%s/%s/%s/%s/%s.png'%(d.path, d.image_id[0], d.image_id[1], d.image_id[2], d.image_id)\n #image_file = data_dir +'/train/%s/%s/%s/%s.png'%(d.image_id[0], d.image_id[1], d.image_id[2], d.image_id)\n image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)\n token = d.sequence\n r = {\n 'index': index,\n 'image_id': d.image_id,\n 'InChI': d.InChI,\n 'formula': d.formula,\n 'd': d,\n 'image': image,\n 'token': token,\n }\n if self.augment is not None:\n r = self.augment(r)\n return r\n\n\ndef null_collate(batch, is_sort_decreasing_length=True):\n collate = defaultdict(list)\n\n if is_sort_decreasing_length: # sort by decreasing length\n sort = np.argsort([-len(r['token']) for r in batch])\n batch = [batch[s] for s in sort]\n\n for r in batch:\n for k, v in r.items():\n collate[k].append(v)\n\n collate['length'] = [len(l) for l in collate['token']]\n\n token = [np.array(t, np.int32) for t in collate['token']]\n token = pad_sequence_to_max_length(token, max_length=max_length, padding_value=STOI[''])\n collate['token'] = torch.from_numpy(token).long()\n\n image = np.stack(collate['image'])\n image = image.astype(np.float32) / 255\n collate['image'] = torch.from_numpy(image).unsqueeze(1).repeat(1, 3, 1, 1)\n\n return collate","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:07.956354Z","iopub.execute_input":"2021-06-03T11:34:07.956747Z","iopub.status.idle":"2021-06-03T11:34:09.648784Z","shell.execute_reply.started":"2021-06-03T11:34:07.956707Z","shell.execute_reply":"2021-06-03T11:34:09.647949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\n\n\n# http://stackoverflow.com/questions/34950201/pycharm-print-end-r-statement-not-working\nclass Logger(object):\n def __init__(self):\n self.terminal = sys.stdout #stdout\n self.file = None\n\n def open(self, file, mode=None):\n if mode is None:\n mode = 'w'\n self.file = open(file, mode)\n\n def write(self, message, is_terminal=1, is_file=1 ):\n if '\\r' in message: is_file=0\n\n if is_terminal == 1:\n self.terminal.write(message)\n self.terminal.flush()\n #time.sleep(1)\n\n if is_file == 1:\n self.file.write(message)\n self.file.flush()\n\n def flush(self):\n # this flush method is needed for python 3 compatibility.\n # this handles the flush command by doing nothing.\n # you might want to specify some extra behavior here.\n pass\n\n\ndef time_to_str(t, mode='min'):\n if mode == 'min':\n t = int(t)/60\n hr = t//60\n min = t % 60\n return '%2d hr %02d min' % (hr, min)\n\n elif mode=='sec':\n t = int(t)\n min = t//60\n sec = t % 60\n return '%2d min %02d sec' % (min, sec)\n\n else:\n raise NotImplementedError\n\n\ndef get_learning_rate(optimizer):\n lr = []\n for param_group in optimizer.param_groups:\n lr += [param_group['lr']]\n\n assert(len(lr) == 1) # we support only one param_group\n lr = lr[0]\n\n return lr","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:09.650085Z","iopub.execute_input":"2021-06-03T11:34:09.650450Z","iopub.status.idle":"2021-06-03T11:34:09.660470Z","shell.execute_reply.started":"2021-06-03T11:34:09.650410Z","shell.execute_reply":"2021-06-03T11:34:09.659628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n\n\nclass Beam:\n\n def __init__(self, beam_size=8, min_length=0, n_top=1, ranker=None,\n start_token_id=2, end_token_id=3):\n self.beam_size = beam_size\n self.min_length = min_length\n self.ranker = ranker\n\n self.end_token_id = end_token_id\n self.top_sentence_ended = False\n\n self.prev_ks = []\n self.next_ys = [torch.LongTensor(beam_size).fill_(start_token_id)] # remove padding\n\n self.current_scores = torch.FloatTensor(beam_size).zero_()\n self.all_scores = []\n\n # The attentions (matrix) for each time.\n self.all_attentions = []\n\n self.finished = []\n\n\n\n # Time and k pair for finished.\n self.finished = []\n self.n_top = n_top\n\n self.ranker = ranker\n\n def advance(self, next_log_probs, current_attention):\n # next_probs : beam_size X vocab_size\n # current_attention: (target_seq_len=1, beam_size, source_seq_len)\n\n vocabulary_size = next_log_probs.size(1)\n # current_beam_size = next_log_probs.size(0)\n\n current_length = len(self.next_ys)\n if current_length < self.min_length:\n for beam_index in range(len(next_log_probs)):\n next_log_probs[beam_index][self.end_token_id] = -1e10\n\n if len(self.prev_ks) > 0:\n beam_scores = next_log_probs + self.current_scores.unsqueeze(1).expand_as(next_log_probs)\n # Don't let EOS have children.\n last_y = self.next_ys[-1]\n for beam_index in range(last_y.size(0)):\n if last_y[beam_index] == self.end_token_id:\n beam_scores[beam_index] = -1e10 # -1e20 raises error when executing\n else:\n beam_scores = next_log_probs[0]\n flat_beam_scores = beam_scores.view(-1)\n top_scores, top_score_ids = flat_beam_scores.topk(k=self.beam_size, dim=0, largest=True, sorted=True)\n\n self.current_scores = top_scores\n self.all_scores.append(self.current_scores)\n\n prev_k = top_score_ids / vocabulary_size # (beam_size, )\n next_y = top_score_ids - prev_k * vocabulary_size # (beam_size, )\n\n self.prev_ks.append(prev_k)\n self.next_ys.append(next_y)\n # for RNN, dim=1 and for transformer, dim=0.\n prev_attention = current_attention.index_select(dim=0, index=prev_k) # (target_seq_len=1, beam_size, source_seq_len)\n self.all_attentions.append(prev_attention)\n\n\n for beam_index, last_token_id in enumerate(next_y):\n if last_token_id == self.end_token_id:\n # skip scoring\n self.finished.append((self.current_scores[beam_index], len(self.next_ys) - 1, beam_index))\n\n if next_y[0] == self.end_token_id:\n self.top_sentence_ended = True\n\n def get_current_state(self):\n \"Get the outputs for the current timestep.\"\n return self.next_ys[-1]\n\n def get_current_origin(self):\n \"Get the backpointers for the current timestep.\"\n return self.prev_ks[-1]\n\n def done(self):\n return self.top_sentence_ended and len(self.finished) >= self.n_top\n\n def get_hypothesis(self, timestep, k):\n hypothesis, attentions = [], []\n for j in range(len(self.prev_ks[:timestep]) - 1, -1, -1):\n hypothesis.append(self.next_ys[j + 1][k])\n # for RNN, [:, k, :], and for trnasformer, [k, :, :]\n attentions.append(self.all_attentions[j][k, :, :])\n k = self.prev_ks[j][k]\n attentions_tensor = torch.stack(attentions[::-1]).squeeze(1) # (timestep, source_seq_len)\n return hypothesis[::-1], attentions_tensor\n\n def sort_finished(self, minimum=None):\n if minimum is not None:\n i = 0\n # Add from beam until we have minimum outputs.\n while len(self.finished) < minimum:\n # global_scores = self.global_scorer.score(self, self.scores)\n # s = global_scores[i]\n s = self.current_scores[i]\n self.finished.append((s, len(self.next_ys) - 1, i))\n i += 1\n\n self.finished = sorted(self.finished, key=lambda a: a[0], reverse=True)\n scores = [sc for sc, _, _ in self.finished]\n ks = [(t, k) for _, t, k in self.finished]\n return scores, ks","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:09.661745Z","iopub.execute_input":"2021-06-03T11:34:09.662330Z","iopub.status.idle":"2021-06-03T11:34:09.684663Z","shell.execute_reply.started":"2021-06-03T11:34:09.662273Z","shell.execute_reply":"2021-06-03T11:34:09.683847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\nimport numpy as np\nfrom typing import Tuple, Dict, Optional\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.rnn import pack_padded_sequence\n\nfrom fairseq.models import *\nfrom fairseq.modules import *\n#from vit import *\n\n# https://arxiv.org/pdf/1411.4555.pdf\n# 'Show and Tell: A Neural Image Caption Generator' - Oriol Vinyals, cvpr-2015\n\nSTOI = {\n '': 190,\n '': 191,\n '': 192,\n #'': 193,\n}\n\n\nclass CNN(nn.Module):\n def __init__(self,type_='vit'):\n super(CNN, self).__init__()\n \n if type_ == 'vit':\n self.e = vit_deit_base_distilled_patch16_384(pretrained=True)\n else:\n self.e = swin_base_patch4_window12_384_in22k(pretrained=True)\n \n for p in self.e.parameters():\n p.requires_grad = True#False\n\n def forward(self, image):\n #batch_size, C, H, W = image.shape\n #x = 2 * image - 1 # ; print('input ', x.size()) \n x = self.e.forward_features(image) ## (bs,img_max_len,image_dim)\n\n return x\n\n\n#https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute\nclass Namespace(object):\n def __init__(self, adict):\n self.__dict__.update(adict)\n\n# ------------------------------------------------------\n# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need\n\nclass PositionEncode1D(nn.Module):\n def __init__(self, dim, max_length):\n super().__init__()\n assert (dim % 2 == 0)\n self.max_length = max_length\n\n d = torch.exp(torch.arange(0., dim, 2) * (-math.log(10000.0) / dim))\n position = torch.arange(0., max_length).unsqueeze(1)\n pos = torch.zeros(1, max_length, dim)\n #pos.require_grad = False\n pos[0, :, 0::2] = torch.sin(position * d)\n pos[0, :, 1::2] = torch.cos(position * d)\n self.register_buffer('pos', pos)\n\n def forward(self, x):\n batch_size, T, dim = x.shape\n x = x + self.pos[:,:T]\n return x\n\n# https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py\n#https://github.com/pytorch/fairseq/issues/568\n# fairseq/fairseq/models/fairseq_encoder.py\n\n# https://github.com/pytorch/fairseq/blob/master/fairseq/modules/transformer_layer.py\nclass TransformerEncode(FairseqEncoder):\n\n def __init__(self, dim, ff_dim, num_head, num_layer):\n super().__init__({})\n #print('my TransformerEncode()')\n\n self.layer = nn.ModuleList([\n TransformerEncoderLayer(Namespace({\n 'encoder_embed_dim': dim,\n 'encoder_attention_heads': num_head,\n 'attention_dropout': 0.1,\n 'dropout': 0.1,\n 'encoder_normalize_before': True,\n 'encoder_ffn_embed_dim': ff_dim,\n })) for i in range(num_layer)\n ])\n self.layer_norm = nn.LayerNorm(dim)\n\n def forward(self, x):# T x B x C\n #print('my TransformerEncode forward()')\n for layer in self.layer:\n x = layer(x)\n x = self.layer_norm(x)\n return x\n\n\n# https://fairseq.readthedocs.io/en/latest/tutorial_simple_lstm.html\n# see https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py\nclass TransformerDecode(FairseqIncrementalDecoder):\n def __init__(self, dim, ff_dim, num_head, num_layer):\n super().__init__({})\n #print('my TransformerDecode()')\n \n #self.layer = LayerDropModuleList(p=0.2)\n #self.layer.extend(\n # [\n # TransformerDecoderLayer(Namespace({\n # 'decoder_embed_dim': dim,\n # 'decoder_attention_heads': num_head,\n # 'attention_dropout': 0.1,\n # 'dropout': 0.1,\n # 'decoder_normalize_before': True,\n # 'decoder_ffn_embed_dim': ff_dim,\n # #'decoder_learned_pos': False,\n # #'cross_self_attention': True,\n # 'activation_fn': 'gelu',\n # })) for i in range(num_layer)\n # ]\n #)\n\n self.layer = nn.ModuleList([\n TransformerDecoderLayer(Namespace({\n 'decoder_embed_dim': dim,\n 'decoder_attention_heads': num_head,\n 'attention_dropout': 0.1,\n 'dropout': 0.1,\n 'decoder_normalize_before': True,\n 'decoder_ffn_embed_dim': ff_dim,\n #'decoder_learned_pos': True,\n #'cross_self_attention': True,\n #'activation-fn': 'gelu',\n })) for i in range(num_layer)\n ])\n self.layer_norm = nn.LayerNorm(dim)\n\n\n def forward(self, x, mem, x_mask):\n #print('my TransformerDecode forward()')\n for layer in self.layer:\n x = layer(x, mem, self_attn_mask=x_mask)[0]\n x = self.layer_norm(x)\n return x # T x B x C\n\n #def forward_one(self, x, mem, incremental_state):\n def forward_one(self,\n x : torch.Tensor,\n mem : torch.Tensor,\n incremental_state : Optional[Dict[str, Dict[str, Optional[torch.Tensor]]]]\n )-> torch.Tensor:\n x = x[-1:]\n for layer in self.layer:\n x = layer(x, mem, incremental_state=incremental_state)[0]\n x = self.layer_norm(x)\n return x\n\n\nclass Net(nn.Module):\n\n def __init__(self,type_,max_length,text_dim,vocab_size,decoder_dim,ff_dim,num_head,num_layer,image_dim):\n super(Net, self).__init__()\n self.type = type_\n \n self.cnn = CNN(type_=type_)\n self.image_encode = nn.Identity()\n # ---\n self.text_pos = PositionEncode1D(text_dim, max_length)\n self.token_embed = nn.Embedding(vocab_size, text_dim)\n self.text_decode = TransformerDecode(decoder_dim, ff_dim, num_head, num_layer)\n\n # ---\n self.logit = nn.Linear(decoder_dim, vocab_size) # (1024, 193)\n #self.logit = nn.Linear(int(decoder_dim/2), vocab_size)\n self.dropout = nn.Dropout(p=0.1)\n\n # ----\n # initialization\n self.token_embed.weight.data.uniform_(-0.1, 0.1)\n self.logit.bias.data.fill_(0)\n self.logit.weight.data.uniform_(-0.1, 0.1)\n\n @torch.jit.unused\n def forward(self, image, token, length):\n device = image.device\n batch_size = len(image)\n # ---\n #print(image.size()) # torch.Size([16, 3, 384, 384])\n image_embed = self.cnn(image)\n #image_embed = F.glu(image_embed)\n #print(image_embed.size()) #torch.Size([16, 144, 1024])\n image_embed = self.image_encode(image_embed).permute(1, 0, 2).contiguous()\n #print(image_embed.size()) #torch.Size([144, 16, 1024])\n \n if self.training:\n probs = torch.empty(token.size()).uniform_(0, 1).to(device)\n probs = torch.where(token < 190, probs, torch.empty(token.size(), dtype=torch.float).fill_(1).to(device))\n token = torch.where(probs > 0.15, token, torch.randint(0, 190, token.size(), dtype=torch.int64).to(device))\n #token = torch.where(probs > 0.33, token, torch.empty(token.size(), dtype=torch.int64).fill_(STOI['']).to(device))\n \n text_embed = self.token_embed(token)\n #print(text_embed.size()) #torch.Size([16, 300, 1024])\n text_embed = self.text_pos(text_embed).permute(1, 0, 2).contiguous()\n #print(text_embed.size()) #torch.Size([300, 16, 1024])\n \n text_mask = np.triu(np.ones((max_length, max_length)), k=1).astype(np.uint8)\n text_mask = torch.autograd.Variable(torch.from_numpy(text_mask) == 1).to(device)\n\n x = self.text_decode(text_embed, image_embed, text_mask)\n x = x.permute(1, 0, 2).contiguous()\n x = self.dropout(x)\n\n logit = self.logit(x) # (193)\n return logit\n\n @torch.jit.export\n def forward_argmax_decode(self, image):\n\n STOI = {\n '': 190,\n '': 191,\n '': 192,\n #'': 193,\n }\n\n image_size = 384\n vocab_size = 193#194\n max_length = 280 #300 # 275\n \n if self.type == 'vit':\n\n image_dim = 768\n text_dim = 768\n decoder_dim = 768\n num_layer = 4 #3\n num_head = 8\n ff_dim = 2048#1024\n \n else:\n image_dim = 1024\n text_dim = 1024\n decoder_dim = 1024\n num_layer = 6 #3\n num_head = 8\n ff_dim = 2048#1024\n \n # ---------------------------------\n device = image.device\n batch_size = len(image)\n\n image_embed = self.cnn(image) #(bs,img_len,image_dim)\n image_embed = self.image_encode(image_embed).permute(1, 0, 2).contiguous() # (img_len,bs,image_dim)\n\n token = torch.full((batch_size, max_length), STOI[''], dtype=torch.long, device=device) # (batch_size,max_len) \n text_pos = self.text_pos.pos #(1,sequence_len,text_dim) torch.zeros(1, max_length, dim)\n token[:, 0] = STOI['']\n\n # -------------------------------------\n eos = STOI['']\n pad = STOI['']\n\n # fast version\n if 1:\n # incremental_state = {}\n incremental_state = torch.jit.annotate(\n Dict[str, Dict[str, Optional[torch.Tensor]]],\n torch.jit.annotate(Dict[str, Dict[str, Optional[torch.Tensor]]], {}),\n )\n for t in range(max_length - 1):\n last_token = token[:, t] # take the whole batch's t'th token\n text_embed = self.token_embed(last_token) #[bs,text_dim] Generate embedding for the t'th token\n text_embed = text_embed + text_pos[:, t] #[bs,text_dim] Combine with pos embed for t'th token\n\n text_embed = text_embed.reshape(1, batch_size, text_dim)\n \n #text_embed ---> 1,bs,text_dim(768)\n #image_embed ---> img_pos_embed,bs,image_im\n x = self.text_decode.forward_one(text_embed, image_embed, incremental_state)\n ## x -----> (1,bs,text_dim)\n x = x.reshape(batch_size, decoder_dim)\n ## x -----> (bs,decoder_dim)\n l = self.logit(x)\n ## l -----> (bs,num_classes)\n k = torch.argmax(l, -1)\n token[:, t + 1] = k\n if ((k == eos) | (k == pad)).all():\n break\n\n predict = token[:, 1:]\n return predict\n\n\nclass EnsembleNet():\n def __init__(self,model_1,model_2):\n self.model_1 = model_1\n self.model_2 = model_2\n \n def forward(self,image):\n \n device = image.device\n batch_size = len(image)\n\n image_embed_1 = self.model_1.cnn(image) #(bs,img_len,image_dim)\n image_embed_1 = self.model_1.image_encode(image_embed_1).permute(1, 0, 2).contiguous() # (img_len,bs,image_dim)\n\n image_embed_2 = self.model_2.cnn(image) #(bs,img_len,image_dim)\n image_embed_2 = self.model_2.image_encode(image_embed_2).permute(1, 0, 2).contiguous() # (img_len,bs,image_dim)\n\n token = torch.full((batch_size, max_length), STOI[''], dtype=torch.long, device=device) # (batch_size,max_len) \n text_pos_1 = self.model_1.text_pos.pos #(1,sequence_len,text_dim) torch.zeros(1, max_length, dim)\n text_pos_2 = self.model_2.text_pos.pos #(1,sequence_len,text_dim) torch.zeros(1, max_length, dim)\n token[:, 0] = STOI['']\n\n# token_2 = torch.full((batch_size, max_length), STOI[''], dtype=torch.long, device=device) # (batch_size,max_len) \n# text_pos_2 = self.model_2.text_pos.pos #(1,sequence_len,text_dim) torch.zeros(1, max_length, dim)\n# token_2[:, 0] = STOI['']\n\n # -------------------------------------\n eos = STOI['']\n pad = STOI['']\n\n # fast version\n if 1:\n # incremental_state = {}\n incremental_state = torch.jit.annotate(\n Dict[str, Dict[str, Optional[torch.Tensor]]],\n torch.jit.annotate(Dict[str, Dict[str, Optional[torch.Tensor]]], {}),\n )\n for t in range(max_length - 1):\n last_token_1 = token[:, t] # take the whole batch's t'th token\n text_embed_1 = self.model_1.token_embed(last_token_1) #[bs,text_dim] Generate embedding for the t'th token\n text_embed_1 = text_embed_1 + text_pos_1[:, t] #[bs,text_dim] Combine with pos embed for t'th token\n\n last_token_2 = token[:, t] # take the whole batch's t'th token\n text_embed_2 = self.model_2.token_embed(last_token_2) #[bs,text_dim] Generate embedding for the t'th token\n text_embed_2 = text_embed_2 + text_pos_2[:, t] #[bs,text_dim] Combine with pos embed for t'th token\n\n text_embed_1 = text_embed_1.reshape(1, batch_size, 768)\n text_embed_2 = text_embed_2.reshape(1, batch_size, 1024)\n\n #text_embed ---> 1,bs,text_dim(768)\n #image_embed ---> img_pos_embed,bs,image_im\n x_1 = self.model_1.text_decode.forward_one(text_embed_1, image_embed_1, incremental_state)\n x_2 = self.model_2.text_decode.forward_one(text_embed_2, image_embed_2, incremental_state)\n ## x -----> (1,bs,text_dim)\n x_1 = x_1.reshape(batch_size,768)\n x_2 = x_2.reshape(batch_size,1024)\n ## x -----> (bs,decoder_dim)\n l_1 = self.model_1.logit(x_1)\n l_2 = self.model_2.logit(x_2)\n\n l = (l_1 + l_2)/2\n ## l -----> (bs,num_classes)\n k = torch.argmax(l, -1)\n token[:, t + 1] = k\n if ((k == eos) | (k == pad)).all():\n break\n\n predict = token[:, 1:]\n return predict\n \n# loss #################################################################\ndef seq_cross_entropy_loss(logit, token, length):\n truth = token[:, 1:]\n L = [l - 1 for l in length]\n logit = pack_padded_sequence(logit, L, batch_first=True).data\n truth = pack_padded_sequence(truth, L, batch_first=True).data\n loss = F.cross_entropy(logit, truth, ignore_index=STOI[''])\n return loss\n\n\n# https://www.aclweb.org/anthology/2020.findings-emnlp.276.pdf\ndef seq_anti_focal_cross_entropy_loss(logit, token, length):\n gamma = 1.0 # {0.5,1.0}\n label_smooth = 0.90\n\n #---\n truth = token[:, 1:]\n L = [l - 1 for l in length]\n logit = pack_padded_sequence(logit, L, batch_first=True).data\n truth = pack_padded_sequence(truth, L, batch_first=True).data\n\n logp = F.log_softmax(logit, -1)\n logp = logp.gather(1, truth.reshape(-1,1)).reshape(-1)\n p = logp.exp()\n\n loss = - ((1 + p) ** gamma)*logp #anti-focal\n loss = loss.mean()\n return loss\n\n\ndef np_loss_cross_entropy(probability, truth):\n batch_size = len(probability)\n truth = truth.reshape(-1)\n p = probability[np.arange(batch_size),truth]\n loss = -np.log(np.clip(p, 1e-6, 1))\n loss = loss.mean()\n return loss","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:09.687277Z","iopub.execute_input":"2021-06-03T11:34:09.687718Z","iopub.status.idle":"2021-06-03T11:34:10.103435Z","shell.execute_reply.started":"2021-06-03T11:34:09.687681Z","shell.execute_reply":"2021-06-03T11:34:10.102633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.environ['CUDA_VISIBLE_DEVICES'] = '0'\nfrom timeit import default_timer as timer\nfrom torch.utils.data import DataLoader, RandomSampler, SequentialSampler\nfrom datetime import datetime","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.104985Z","iopub.execute_input":"2021-06-03T11:34:10.105342Z","iopub.status.idle":"2021-06-03T11:34:10.111037Z","shell.execute_reply.started":"2021-06-03T11:34:10.105306Z","shell.execute_reply":"2021-06-03T11:34:10.109422Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.cuda.amp as amp\nAmpNet = Net\nIDENTIFIER = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')\n\nfold = 3\ninitial_checkpoint_vit = '../input/bms-predictions/deit_vit384_01080000_model.pth'\ninitial_checkpoint_swin = '../input/bms-predictions/swin_00800000_model.pth'\nmode = 'remote'","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.112422Z","iopub.execute_input":"2021-06-03T11:34:10.112996Z","iopub.status.idle":"2021-06-03T11:34:10.121949Z","shell.execute_reply.started":"2021-06-03T11:34:10.112937Z","shell.execute_reply":"2021-06-03T11:34:10.121065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# start here ! ###################################################################################\ndef fast_remote_unrotate_augment(r):\n image = r['image']\n index = r['index']\n h, w = image.shape\n\n if h > w:\n image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)\n\n image = cv2.resize(image, dsize=(image_size,image_size), interpolation=cv2.INTER_LINEAR)\n assert (image_size == 384)\n image = image.astype(np.float16) / 255\n image = torch.from_numpy(image).unsqueeze(0).repeat(3,1,1)\n\n r = {}\n r['image'] = image\n return r","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.123237Z","iopub.execute_input":"2021-06-03T11:34:10.123631Z","iopub.status.idle":"2021-06-03T11:34:10.132528Z","shell.execute_reply.started":"2021-06-03T11:34:10.123593Z","shell.execute_reply":"2021-06-03T11:34:10.131800Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def do_predict(net, tokenizer, valid_loader):\n\n text = []\n\n start_timer = timer()\n valid_num = 0\n for t, batch in enumerate(valid_loader):\n batch_size = len(batch['image'])\n image = batch['image'].cuda()\n\n with torch.no_grad():\n with amp.autocast():\n k = net.forward(image)\n\n # token = batch['token'].cuda()\n # length = batch['length']\n # logit = data_parallel(net,(image, token, length))\n # k = logit.argmax(-1)\n\n k = k.data.cpu().numpy()\n k = tokenizer.predict_to_inchi(k)\n text.extend(k)\n\n valid_num += batch_size\n print('\\r %8d / %d %s' % (valid_num, len(valid_loader.dataset), time_to_str(timer() - start_timer, 'sec')),\n end='', flush=True)\n\n assert(valid_num == len(valid_loader.dataset))\n print('')\n return text","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.133932Z","iopub.execute_input":"2021-06-03T11:34:10.134382Z","iopub.status.idle":"2021-06-03T11:34:10.144784Z","shell.execute_reply.started":"2021-06-03T11:34:10.134343Z","shell.execute_reply":"2021-06-03T11:34:10.143867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_submit():\n gpu_no = int(os.environ['CUDA_VISIBLE_DEVICES'])\n\n is_norm_ichi = False #True\n\n if 1:\n ## setup ----------------------------------------\n submit_dir = 'valid/%s-gpu%d'%(mode, gpu_no)\n os.makedirs(submit_dir, exist_ok=True)\n\n log = Logger()\n log.open('log.submit.txt', mode='a')\n log.write('\\n--- [START %s] %s\\n\\n' % (IDENTIFIER, '-' * 64))\n log.write('is_norm_ichi = %s\\n' % is_norm_ichi)\n log.write('\\n')\n\n #\n ## dataset ------------------------------------\n tokenizer = load_tokenizer()\n\n if 'remote' in mode: # 1_616_107\n df_valid = make_fold('test')\n if gpu_no==0:\n df_valid = df_valid[:400_000]\n if gpu_no==1:\n df_valid = df_valid[400_000:800_000]\n if gpu_no==2:\n df_valid = df_valid[800_000:1200_000]\n if gpu_no==3: \n df_valid = df_valid[1200_000:]\n\n if 'local' in mode: # 484_837\n df_train, df_valid = make_fold('train-%d' % fold)\n\n df_valid = df_valid.sort_values('length').reset_index(drop=True)\n valid_dataset = BmsDataset(df_valid, tokenizer, augment=fast_remote_unrotate_augment)\n valid_loader = DataLoader(\n valid_dataset,\n sampler=SequentialSampler(valid_dataset),\n batch_size=6, #128\n drop_last=False,\n num_workers=4,\n pin_memory=True,\n )\n log.write('mode : %s\\n'%(mode))\n log.write('valid_dataset : \\n%s\\n'%(valid_dataset))\n\n ## net ----------------------------------------\n if 1:\n tokenizer = load_tokenizer()\n net_1 = AmpNet(type_='vit',max_length=280,text_dim=768,vocab_size=193,decoder_dim=768,ff_dim=2048,num_head=8,num_layer=4,image_dim = 768).cuda()\n net_1.load_state_dict(torch.load(initial_checkpoint_vit)['state_dict'], strict=True)\n net_1 = torch.jit.script(net_1)\n net_1.eval()\n \n net_2 = AmpNet(type_='swin',max_length=280,text_dim=1024,vocab_size=193,decoder_dim=1024,ff_dim=2048,num_head=8,num_layer=6,image_dim = 1024).cuda()\n net_2.load_state_dict(torch.load(initial_checkpoint_swin)['state_dict'], strict=True)\n net_2 = torch.jit.script(net_2)\n net_2.eval()\n \n net = EnsembleNet(net_1,net_2)\n\n #---\n start_timer = timer()\n predict = do_predict(net, tokenizer, valid_loader)\n log.write('time %s \\n' % time_to_str(timer() - start_timer, 'min'))\n else:\n pass\n #----\n if is_norm_ichi:\n predict = [normalize_inchi(t) for t in predict] #\n\n df_submit = pd.DataFrame()\n df_submit.loc[:, 'image_id'] = df_valid.image_id.values\n df_submit.loc[:, 'InChI'] = predict #\n df_submit.to_csv(submit_dir + '/submit.csv', index=False)\n\n log.write('submit_dir : %s\\n' % (submit_dir))\n log.write('initial_checkpoint : %s\\n' % (initial_checkpoint))\n log.write('df_submit : %s\\n' % str(df_submit.shape))\n log.write('%s\\n' % str(df_submit))\n log.write('\\n')\n\n if 'local' in mode:\n truth = df_valid['InChI'].values.tolist()\n lb_score = compute_lb_score(predict, truth)\n #print(lb_score)\n log.write('lb_score = %f\\n'%lb_score.mean())\n log.write('is_norm_ichi = %s\\n' % is_norm_ichi)\n log.write('\\n')\n\n if 1:\n df_eval = df_submit.copy()\n df_eval.loc[:,'truth']=truth\n df_eval.loc[:,'lb_score']=lb_score\n df_eval.loc[:,'length'] = df_valid['length']\n df_eval.to_csv(submit_dir + '/df_eval.csv', index=False)\n # df_valid.to_csv(submit_dir + '/df_valid', index=False)\n\n #exit(0)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.146227Z","iopub.execute_input":"2021-06-03T11:34:10.146687Z","iopub.status.idle":"2021-06-03T11:34:10.167223Z","shell.execute_reply.started":"2021-06-03T11:34:10.146647Z","shell.execute_reply":"2021-06-03T11:34:10.166421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"####################################################### main #################################################################\nif __name__ == '__main__':\n run_submit()","metadata":{"execution":{"iopub.status.busy":"2021-06-03T11:34:10.168599Z","iopub.execute_input":"2021-06-03T11:34:10.169063Z","iopub.status.idle":"2021-06-03T11:35:02.321914Z","shell.execute_reply.started":"2021-06-03T11:34:10.169029Z","shell.execute_reply":"2021-06-03T11:35:02.320060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}