{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install geffnet\n!pip install git+https://github.com/pabloppp/pytorch-tools -U\n!pip install thop ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data=pd.read_csv('/kaggle/input/state-farm-distracted-driver-detection/driver_imgs_list.csv')\ndata.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data['classname'].value_counts()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_label_dict={}\nfor i in range(data.shape[0]):\n    info=data.iloc[i]\n    image_name=info['img']\n    label=info['classname']\n    image_label_dict[image_name]=label\n    \n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from fastai.vision.all import get_image_files\nimport random\nfrom sklearn.model_selection import KFold,StratifiedKFold\nnames=get_image_files('/kaggle/input/state-farm-distracted-driver-detection/imgs/train')\nrandom.shuffle(names)\nX=[]\nY=[]\nfor i in range(len(names)):\n    name=str(names[i])\n    name=name.split('/')[-1]\n    label=image_label_dict[name]\n    X.append([i])\n    Y.append (label)\n    \nsfolder = StratifiedKFold(n_splits=5,random_state=0,shuffle=True)\n\ntrain,test=next(sfolder.split(X,Y))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_name=[]\ntest_name=[]\ntrain_label=[]\ntest_label=[]\nfor i in train:\n    index=X[i][0]\n    train_name.append(names[index])\n    train_label.append(Y[i])\nfor i in test:\n    test_name.append(names[index])\n    test_label.append(Y[i])\n    \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.utils.data as Data\nfrom PIL import Image\nimport torchvision.transforms as T\nimport torch.nn.functional as F\ntrain_transform=T.Compose([\n       T.Resize([384,384]),\n       T.RandomHorizontalFlip(),\n       T.Pad(10),\n       T.RandomCrop([384,384]),\n       T.ColorJitter(),\n       T.RandomAffine(degrees=10),\n       T.ToTensor(),\n       T.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])\n]\n)\ntest_transform=T.Compose([\n    T.Resize([384,384]),\n    T.ToTensor(),\n    T.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])\n])\n\n\nclass Dataset(Data.Dataset):\n    def __init__(self,name,label,transform):\n        super(Dataset, self).__init__()\n        self.name=name\n        self.label=label\n        self.transform=transform\n\n    def __getitem__(self, index):\n        path=str(self.name[index])\n        label=self.label[index]\n        label=int(label.strip()[-1])\n        image=Image.open(path)\n        image=self.transform(image)\n        return image,label\n\n    def __len__(self):\n        return len(self.name)\n\nclass GhostBottleneck(nn.Module):\n    \"\"\" Ghost bottleneck w/ optional SE\"\"\"\n\n    def __init__(self, in_chs, mid_chs, out_chs, dw_kernel_size=3,\n                 stride=1, act_layer=nn.ReLU, se_ratio=0.):\n        super(GhostBottleneck, self).__init__()\n        has_se = se_ratio is not None and se_ratio > 0.\n        self.stride = stride\n\n        # Point-wise expansion\n        self.ghost1 = GhostModule(in_chs, mid_chs, relu=True)\n\n        # Depth-wise convolution\n        if self.stride > 1:\n            self.conv_dw = nn.Conv2d(mid_chs, mid_chs, dw_kernel_size, stride=stride,\n                                     padding=(dw_kernel_size - 1) // 2,\n                                     groups=mid_chs, bias=False)\n            self.bn_dw = nn.BatchNorm2d(mid_chs)\n\n        # Squeeze-and-excitation\n        if has_se:\n            self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio)\n        else:\n            self.se = None\n\n        # Point-wise linear projection\n        self.ghost2 = GhostModule(mid_chs, out_chs, relu=False)\n\n        # shortcut\n        if (in_chs == out_chs and self.stride == 1):\n            self.shortcut = nn.Sequential()\n        else:\n            self.shortcut = nn.Sequential(\n                nn.Conv2d(in_chs, in_chs, dw_kernel_size, stride=stride,\n                          padding=(dw_kernel_size - 1) // 2, groups=in_chs, bias=False),\n                nn.BatchNorm2d(in_chs),\n                nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False),\n                nn.BatchNorm2d(out_chs),\n            )\n\n    def forward(self, x):\n        residual = x\n\n        # 1st ghost bottleneck\n        x = self.ghost1(x)\n\n        # Depth-wise convolution\n        if self.stride > 1:\n            x = self.conv_dw(x)\n            x = self.bn_dw(x)\n\n        # Squeeze-and-excitation\n        if self.se is not None:\n            x = self.se(x)\n\n        # 2nd ghost bottleneck\n        x = self.ghost2(x)\n\n        x += self.shortcut(residual)\n        return x\n\ndef _make_divisible(v, divisor, min_value=None):\n    \"\"\"\n    This function is taken from the original tf repo.\n    It ensures that all layers have a channel number that is divisible by 8\n    It can be seen here:\n    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py\n    \"\"\"\n    if min_value is None:\n        min_value = divisor\n    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)\n    # Make sure that round down does not go down by more than 10%.\n    if new_v < 0.9 * v:\n        new_v += divisor\n    return new_v\n\n\ndef hard_sigmoid(x, inplace: bool = False):\n    if inplace:\n        return x.add_(3.).clamp_(0., 6.).div_(6.)\n    else:\n        return F.relu6(x + 3.) / 6.\n\nimport math\nclass SqueezeExcite(nn.Module):\n    def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None,\n                 act_layer=nn.ReLU, gate_fn=hard_sigmoid, divisor=4, **_):\n        super(SqueezeExcite, self).__init__()\n        self.gate_fn = gate_fn\n        reduced_chs = _make_divisible((reduced_base_chs or in_chs) * se_ratio, divisor)\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)\n        self.act1 = act_layer(inplace=True)\n        self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)\n\n    def forward(self, x):\n        x_se = self.avg_pool(x)\n        x_se = self.conv_reduce(x_se)\n        x_se = self.act1(x_se)\n        x_se = self.conv_expand(x_se)\n        x = x * self.gate_fn(x_se)\n        return x\n\n\nclass ConvBnAct(nn.Module):\n    def __init__(self, in_chs, out_chs, kernel_size,\n                 stride=1, act_layer=nn.ReLU):\n        super(ConvBnAct, self).__init__()\n        self.conv = nn.Conv2d(in_chs, out_chs, kernel_size, stride, kernel_size // 2, bias=False)\n        self.bn1 = nn.BatchNorm2d(out_chs)\n        self.act1 = act_layer(inplace=True)\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        return x\n\n\nclass GhostModule(nn.Module):\n    def __init__(self, inp, oup, kernel_size=1, ratio=2, dw_size=3, stride=1, relu=True):\n        super(GhostModule, self).__init__()\n        self.oup = oup\n        init_channels = math.ceil(oup / ratio)\n        new_channels = init_channels * (ratio - 1)\n\n        self.primary_conv = nn.Sequential(\n            nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size // 2, bias=False),\n            nn.BatchNorm2d(init_channels),\n            nn.ReLU(inplace=True) if relu else nn.Sequential(),\n        )\n\n        self.cheap_operation = nn.Sequential(\n            nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size // 2, groups=init_channels, bias=False),\n            nn.BatchNorm2d(new_channels),\n            nn.ReLU(inplace=True) if relu else nn.Sequential(),\n        )\n\n    def forward(self, x):\n        x1 = self.primary_conv(x)\n        x2 = self.cheap_operation(x1)\n        out = torch.cat([x1, x2], dim=1)\n        return out[:, :self.oup, :, :]\n\n\nclass GhostNet(nn.Module):\n    def __init__(self, cfgs, num_classes=1000, width=1.0, dropout=0.2):\n        super(GhostNet, self).__init__()\n        # setting of inverted residual blocks\n        self.cfgs = cfgs\n        self.dropout = dropout\n\n        # building first layer\n        output_channel = _make_divisible(16 * width, 4)\n        self.conv_stem = nn.Conv2d(3, output_channel, 3, 2, 1, bias=False)\n        self.bn1 = nn.BatchNorm2d(output_channel)\n        self.act1 = nn.ReLU(inplace=True)\n        input_channel = output_channel\n\n        # building inverted residual blocks\n        stages = []\n        block = GhostBottleneck\n        for cfg in self.cfgs:\n            layers = []\n            for k, exp_size, c, se_ratio, s in cfg:\n                output_channel = _make_divisible(c * width, 4)\n                hidden_channel = _make_divisible(exp_size * width, 4)\n                layers.append(block(input_channel, hidden_channel, output_channel, k, s,\n                                    se_ratio=se_ratio))\n                input_channel = output_channel\n            stages.append(nn.Sequential(*layers))\n\n        output_channel = _make_divisible(exp_size * width, 4)\n        stages.append(nn.Sequential(ConvBnAct(input_channel, output_channel, 1)))\n        input_channel = output_channel\n\n        self.blocks = nn.Sequential(*stages)\n        self.conv_final=nn.Sequential(\n            nn.Conv2d(960,2048,kernel_size=1),\n            nn.BatchNorm2d(2048),\n            nn.ReLU()\n        )\n        # building last several layers\n        output_channel = 1280\n\n\n    def forward(self, x):\n        x = self.conv_stem(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        x = self.blocks(x)\n        x=self.conv_final(x)\n\n        return x\n\ndef ghostnet(**kwargs):\n    \"\"\"\n    Constructs a GhostNet model\n    \"\"\"\n    cfgs = [\n        # k, t, c, SE, s\n        # stage1\n        [[3,  16,  16, 0, 1]],\n        # stage2\n        [[3,  48,  24, 0, 2]],\n        [[3,  72,  24, 0, 1]],\n        # stage3\n        [[5,  72,  40, 0.25, 2]],\n        [[5, 120,  40, 0.25, 1]],\n        # stage4\n        [[3, 240,  80, 0, 2]],\n        [[3, 200,  80, 0, 1],\n         [3, 184,  80, 0, 1],\n         [3, 184,  80, 0, 1],\n         [3, 480, 112, 0.25, 1],\n         [3, 672, 112, 0.25, 1]\n        ],\n        # stage5\n        [[5, 672, 160, 0.25, 2]],\n        [[5, 960, 160, 0, 1],\n         [5, 960, 160, 0.25, 1],\n         [5, 960, 160, 0, 1],\n         [5, 960, 160, 0.25, 1]\n        ]\n    ]\n    return GhostNet(cfgs, **kwargs)\nclass GeneralizedMeanPooling(nn.Module):\n    r\"\"\"Applies a 2D power-average adaptive pooling over an input signal composed of several input planes.\n    The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)`\n        - At p = infinity, one gets Max Pooling\n        - At p = 1, one gets Average Pooling\n    The output is of size H x W, for any input size.\n    The number of output features is equal to the number of input planes.\n    Args:\n        output_size: the target output size of the image of the form H x W.\n                     Can be a tuple (H, W) or a single H for a square image H x H\n                     H and W can be either a ``int``, or ``None`` which means the size will\n                     be the same as that of the input.\n    \"\"\"\n\n    def __init__(self, norm, output_size=1, eps=1e-6):\n        super(GeneralizedMeanPooling, self).__init__()\n        assert norm > 0\n        self.p = float(norm)\n        self.output_size = output_size\n        self.eps = eps\n\n    def forward(self, x):\n        x = x.clamp(min=self.eps).pow(self.p)\n        return torch.nn.functional.adaptive_avg_pool2d(x, self.output_size).pow(1. / self.p)\n\n    def __repr__(self):\n        return self.__class__.__name__ + '(' \\\n               + str(self.p) + ', ' \\\n               + 'output_size=' + str(self.output_size) + ')'\n\n\nclass GeneralizedMeanPoolingP(GeneralizedMeanPooling):\n    \"\"\" Same, but norm is trainable\n    \"\"\"\n\n    def __init__(self, norm=3, output_size=1, eps=1e-6):\n        super(GeneralizedMeanPoolingP, self).__init__(norm, output_size, eps)\n        self.p = nn.Parameter(torch.ones(1) * norm)\n\n\nclass ConvBlock(nn.Module):\n    \"\"\"Basic convolutional block.\n\n    convolution + batch normalization + relu.\n    Args:\n        in_c (int): number of input channels.\n        out_c (int): number of output channels.\n        k (int or tuple): kernel size.\n        s (int or tuple): stride.\n        p (int or tuple): padding.\n    \"\"\"\n\n    def __init__(self, in_c, out_c, k, s=1, p=0):\n        super(ConvBlock, self).__init__()\n        self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)\n        self.bn = nn.BatchNorm2d(out_c)\n\n    def forward(self, x):\n        return F.relu(self.bn(self.conv(x)))\n\n\nclass InceptionA(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super(InceptionA, self).__init__()\n        mid_channels = out_channels // 4\n\n        self.stream1 = nn.Sequential(\n            ConvBlock(in_channels, mid_channels, 1),\n            ConvBlock(mid_channels, mid_channels, 3, p=1),\n        )\n        self.stream2 = nn.Sequential(\n            ConvBlock(in_channels, mid_channels, 1),\n            ConvBlock(mid_channels, mid_channels, 3, p=1),\n        )\n        self.stream3 = nn.Sequential(\n            ConvBlock(in_channels, mid_channels, 1),\n            ConvBlock(mid_channels, mid_channels, 3, p=1),\n        )\n        self.stream4 = nn.Sequential(\n            nn.AvgPool2d(3, stride=1, padding=1),\n            ConvBlock(in_channels, mid_channels, 1),\n        )\n\n    def forward(self, x):\n        s1 = self.stream1(x)\n        s2 = self.stream2(x)\n        s3 = self.stream3(x)\n        s4 = self.stream4(x)\n        y = torch.cat([s1, s2, s3, s4], dim=1)\n        return y\n\n\nclass InceptionB(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super(InceptionB, self).__init__()\n        mid_channels = out_channels // 4\n\n        self.stream1 = nn.Sequential(\n            ConvBlock(in_channels, mid_channels, 1),\n            ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),\n        )\n        self.stream2 = nn.Sequential(\n            ConvBlock(in_channels, mid_channels, 1),\n            ConvBlock(mid_channels, mid_channels, 3, p=1),\n            ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),\n        )\n        self.stream3 = nn.Sequential(\n            nn.MaxPool2d(3, stride=2, padding=1),\n            ConvBlock(in_channels, mid_channels * 2, 1),\n        )\n\n    def forward(self, x):\n        s1 = self.stream1(x)\n        s2 = self.stream2(x)\n        s3 = self.stream3(x)\n        y = torch.cat([s1, s2, s3], dim=1)\n        return y\n\n\nclass SpatialAttn(nn.Module):\n    \"\"\"Spatial Attention (Sec. 3.1.I.1)\"\"\"\n\n    def __init__(self):\n        super(SpatialAttn, self).__init__()\n        self.conv1 = ConvBlock(1, 1, 3, s=2, p=1)\n        self.conv2 = ConvBlock(1, 1, 1)\n\n    def forward(self, x):\n        # global cross-channel averaging\n        x = x.mean(1, keepdim=True)\n        # 3-by-3 conv\n        x = self.conv1(x)\n        # bilinear resizing\n        x = F.upsample(\n            x, (x.size(2) * 2, x.size(3) * 2),\n            mode='bilinear',\n            align_corners=True\n        )\n        # scaling conv\n        x = self.conv2(x)\n        return x\n\n\nclass ChannelAttn(nn.Module):\n    \"\"\"Channel Attention (Sec. 3.1.I.2)\"\"\"\n\n    def __init__(self, in_channels, reduction_rate=16):\n        super(ChannelAttn, self).__init__()\n        assert in_channels % reduction_rate == 0\n        self.conv1 = ConvBlock(in_channels, in_channels // reduction_rate, 1)\n        self.conv2 = ConvBlock(in_channels // reduction_rate, in_channels, 1)\n\n    def forward(self, x):\n        # squeeze operation (global average pooling)\n        x = F.avg_pool2d(x, x.size()[2:])\n        # excitation operation (2 conv layers)\n        x = self.conv1(x)\n        x = self.conv2(x)\n        return x\n\n\nclass SoftAttn(nn.Module):\n    \"\"\"Soft Attention (Sec. 3.1.I)\n\n    Aim: Spatial Attention + Channel Attention\n\n    Output: attention maps with shape identical to input.\n    \"\"\"\n\n    def __init__(self, in_channels):\n        super(SoftAttn, self).__init__()\n        self.spatial_attn = SpatialAttn()\n        self.channel_attn = ChannelAttn(in_channels)\n        self.conv = ConvBlock(in_channels, in_channels, 1)\n\n    def forward(self, x):\n        y_spatial = self.spatial_attn(x)\n        y_channel = self.channel_attn(x)\n        y = y_spatial * y_channel\n        y = torch.sigmoid(self.conv(y))\n        return y\n\n\nclass HardAttn(nn.Module):\n    \"\"\"Hard Attention (Sec. 3.1.II)\"\"\"\n\n    def __init__(self, in_channels):\n        super(HardAttn, self).__init__()\n        self.fc = nn.Linear(in_channels, 4 * 2)\n        self.init_params()\n\n    def init_params(self):\n        self.fc.weight.data.zero_()\n        self.fc.bias.data.copy_(\n            torch.tensor(\n                [0, -0.75, 0, -0.25, 0, 0.25, 0, 0.75], dtype=torch.float\n            )\n        )\n\n    def forward(self, x):\n        # squeeze operation (global average pooling)\n        x = F.avg_pool2d(x, x.size()[2:]).view(x.size(0), x.size(1))\n        # predict transformation parameters\n        theta = torch.tanh(self.fc(x))\n        theta = theta.view(-1, 4, 2)\n        return theta\n\n\nclass HarmAttn(nn.Module):\n    \"\"\"Harmonious Attention (Sec. 3.1)\"\"\"\n\n    def __init__(self, in_channels):\n        super(HarmAttn, self).__init__()\n        self.soft_attn = SoftAttn(in_channels)\n        self.hard_attn = HardAttn(in_channels)\n\n    def forward(self, x):\n        y_soft_attn = self.soft_attn(x)\n        theta = self.hard_attn(x)\n        return y_soft_attn, theta\n\n\nclass HACNN(nn.Module):\n    \"\"\"Harmonious Attention Convolutional Neural Network.\n    Reference:\n        Li et al. Harmonious Attention Network for Person Re-identification. CVPR 2018.\n    Public keys:\n        - ``hacnn``: HACNN.\n    \"\"\"\n\n    # Args:\n    #    num_classes (int): number of classes to predict\n    #    nchannels (list): number of channels AFTER concatenation\n    #    feat_dim (int): feature dimension for a single stream\n    #    learn_region (bool): whether to learn region features (i.e. local branch)\n\n    def __init__(\n            self,\n            num_classes,\n            loss='softmax',\n            nchannels=[128, 256, 384],\n            feat_dim=512,\n            learn_region=True,\n            use_gpu=True,\n            **kwargs\n    ):\n        super(HACNN, self).__init__()\n        self.loss = loss\n        self.learn_region = learn_region\n        self.use_gpu = use_gpu\n\n        self.conv = ConvBlock(3, 32, 3, s=2, p=1)\n\n        # Construct Inception + HarmAttn blocks\n        # ============== Block 1 ==============\n        self.inception1 = nn.Sequential(\n            InceptionA(32, nchannels[0]),\n            InceptionB(nchannels[0], nchannels[0]),\n        )\n        self.ha1 = HarmAttn(nchannels[0])\n\n        # ============== Block 2 ==============\n        self.inception2 = nn.Sequential(\n            InceptionA(nchannels[0], nchannels[1]),\n            InceptionB(nchannels[1], nchannels[1]),\n        )\n        self.ha2 = HarmAttn(nchannels[1])\n\n        # ============== Block 3 ==============\n        self.inception3 = nn.Sequential(\n            InceptionA(nchannels[1], nchannels[2]),\n            InceptionB(nchannels[2], nchannels[2]),\n        )\n        self.ha3 = HarmAttn(nchannels[2])\n\n        self.fc_global = nn.Sequential(\n            nn.Linear(nchannels[2], feat_dim),\n            nn.BatchNorm1d(feat_dim),\n            nn.ReLU(),\n        )\n        self.classifier_global = nn.Linear(feat_dim, num_classes)\n\n        if self.learn_region:\n            self.init_scale_factors()\n            self.local_conv1 = InceptionB(32, nchannels[0])\n            self.local_conv2 = InceptionB(nchannels[0], nchannels[1])\n            self.local_conv3 = InceptionB(nchannels[1], nchannels[2])\n            self.fc_local = nn.Sequential(\n                nn.Linear(nchannels[2] * 4, feat_dim),\n                nn.BatchNorm1d(feat_dim),\n                nn.ReLU(),\n            )\n            self.classifier_local = nn.Linear(feat_dim, num_classes)\n            self.feat_dim = feat_dim * 2\n        else:\n            self.feat_dim = feat_dim\n\n    def init_scale_factors(self):\n        # initialize scale factors (s_w, s_h) for four regions\n        self.scale_factors = []\n        self.scale_factors.append(\n            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n        )\n        self.scale_factors.append(\n            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n        )\n        self.scale_factors.append(\n            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n        )\n        self.scale_factors.append(\n            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n        )\n\n    def stn(self, x, theta):\n        \"\"\"Performs spatial transform\n\n        x: (batch, channel, height, width)\n        theta: (batch, 2, 3)\n        \"\"\"\n        grid = F.affine_grid(theta, x.size())\n        x = F.grid_sample(x, grid)\n        return x\n\n    def transform_theta(self, theta_i, region_idx):\n        \"\"\"Transforms theta to include (s_w, s_h), resulting in (batch, 2, 3)\"\"\"\n        scale_factors = self.scale_factors[region_idx]\n        theta = torch.zeros(theta_i.size(0), 2, 3)\n        theta[:, :, :2] = scale_factors\n        theta[:, :, -1] = theta_i\n        if self.use_gpu:\n            theta = theta.cuda()\n        return theta\n\n    def forward(self, x):\n        x = self.conv(x)\n\n        # ============== Block 1 ==============\n        # global branch\n        x1 = self.inception1(x)\n        x1_attn, x1_theta = self.ha1(x1)\n        x1_out = x1 * x1_attn\n        # local branch\n        if self.learn_region:\n            x1_local_list = []\n            for region_idx in range(4):\n                x1_theta_i = x1_theta[:, region_idx, :]\n                x1_theta_i = self.transform_theta(x1_theta_i, region_idx)\n                x1_trans_i = self.stn(x, x1_theta_i)\n                x1_trans_i = F.upsample(\n                    x1_trans_i, (24, 28), mode='bilinear', align_corners=True\n                )\n                x1_local_i = self.local_conv1(x1_trans_i)\n                x1_local_list.append(x1_local_i)\n\n        # ============== Block 2 ==============\n        # Block 2\n        # global branch\n        x2 = self.inception2(x1_out)\n        x2_attn, x2_theta = self.ha2(x2)\n        x2_out = x2 * x2_attn\n        # local branch\n        if self.learn_region:\n            x2_local_list = []\n            for region_idx in range(4):\n                x2_theta_i = x2_theta[:, region_idx, :]\n                x2_theta_i = self.transform_theta(x2_theta_i, region_idx)\n                x2_trans_i = self.stn(x1_out, x2_theta_i)\n                x2_trans_i = F.upsample(\n                    x2_trans_i, (12, 14), mode='bilinear', align_corners=True\n                )\n                x2_local_i = x2_trans_i + x1_local_list[region_idx]\n                x2_local_i = self.local_conv2(x2_local_i)\n                x2_local_list.append(x2_local_i)\n\n        # ============== Block 3 ==============\n        # Block 3\n        # global branch\n        x3 = self.inception3(x2_out)\n        x3_attn, x3_theta = self.ha3(x3)\n        x3_out = x3 * x3_attn\n        # local branch\n        if self.learn_region:\n            x3_local_list = []\n            for region_idx in range(4):\n                x3_theta_i = x3_theta[:, region_idx, :]\n                x3_theta_i = self.transform_theta(x3_theta_i, region_idx)\n                x3_trans_i = self.stn(x2_out, x3_theta_i)\n                x3_trans_i = F.upsample(\n                    x3_trans_i, (6, 7), mode='bilinear', align_corners=True\n                )\n                x3_local_i = x3_trans_i + x2_local_list[region_idx]\n                x3_local_i = self.local_conv3(x3_local_i)\n                x3_local_list.append(x3_local_i)\n\n        global_feat=x3_out\n        x_local_list=[]\n        for region_idx in range(4):\n            x_local_i = x3_local_list[region_idx]\n            x_local_list.append(x_local_i)\n        x_local = torch.cat(x_local_list, 1)\n        feature=torch.cat([global_feat, x_local], 1)\n\n        return feature\n\n\n\nfrom torchvision.models import resnet50\nfrom geffnet import efficientnet_b3\ndef build_backbone(name):\n    if name=='resnet50':\n        model = resnet50(True)\n        stride = 1\n        model.layer4[0].downsample[0].stride = stride\n        model.layer4[0].conv2.stride = stride\n        base = nn.Sequential(\n            model.conv1,\n            model.bn1,\n            model.maxpool,\n            model.layer1,\n            model.layer2,\n            model.layer3,\n            model.layer4\n        )\n        return base, 2048\n    elif name=='eff_b3':\n        model = efficientnet_b3(True)\n        base = nn.Sequential(\n            model.conv_stem,\n            model.bn1,\n            model.act1,\n            model.blocks,\n            model.conv_head,\n            model.bn2,\n            model.act2\n        )\n        return base, 1536\n    elif name=='ghost':\n        base=ghostnet()\n        return base,2048\n    elif name=='hanet':\n        base=HACNN()\n        return base,384*5\n\n\n\ndef weights_init_kaiming(m):\n    classname = m.__class__.__name__\n    if classname.find('Linear') != -1:\n        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')\n        nn.init.constant_(m.bias, 0.0)\n    elif classname.find('Conv') != -1:\n        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')\n        if m.bias is not None:\n            nn.init.constant_(m.bias, 0.0)\n    elif classname.find('BatchNorm') != -1:\n        if m.affine:\n            nn.init.constant_(m.weight, 1.0)\n            nn.init.constant_(m.bias, 0.0)\n\ndef weights_init_classifier(m):\n    classname = m.__class__.__name__\n    if classname.find('Linear') != -1:\n        nn.init.normal_(m.weight, std=0.001)\n        if m.bias:\n            nn.init.constant_(m.bias, 0.0)\n\n\n\nclass Model(nn.Module):\n    def __init__(self,name,num_class):\n        super(Model, self).__init__()\n        self.backbone,feat=build_backbone(name)\n        self.pool=GeneralizedMeanPoolingP()\n        self.bottleneck = nn.BatchNorm1d(feat)\n        self.bottleneck.bias.requires_grad_(False)  # no shift\n        self.classifier = nn.Linear(feat, num_class, bias=False)\n\n        self.bottleneck.apply(weights_init_kaiming)\n        self.classifier.apply(weights_init_classifier)\n\n    def forward(self,x):\n        x=self.backbone(x)\n        x=self.pool(x)\n        global_feature=x.view(x.shape[0],-1)\n        feature=self.bottleneck(global_feature)\n        logit=self.classifier(feature)\n        if self.training:\n            return global_feature,logit\n        else:\n            return logit\n\n\nclass Flat(torch.optim.lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, epochs, anneal_start=0.65, last_epoch=-1):\n        self.epochs = epochs\n        self.start = anneal_start\n        super(Flat, self).__init__(optimizer, last_epoch)\n\n    def get_lr(self):\n        if self.last_epoch < self.epochs * self.start:\n            return [base_lr for base_lr in self.base_lrs]\n        else:\n            return [\n                base_lr * (1 + math.cos(math.pi * self.last_epoch / 5)) / 2\n                for base_lr in self.base_lrs\n            ]\n\n\nimport torch.nn as nn\n\n\nclass CrossEntropyLabelSmoothLoss(nn.Module):\n    \"\"\"Cross entropy loss with label smoothing regularizer.\n    Reference:\n    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.\n    Equation: y = (1 - epsilon) * y + epsilon / K.\n    Args:\n        num_classes (int): number of classes.\n        epsilon (float): weight.\n    \"\"\"\n\n    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):\n        super(CrossEntropyLabelSmoothLoss, self).__init__()\n        self.num_classes = num_classes\n        self.epsilon = epsilon\n        self.use_gpu = use_gpu\n        self.logsoftmax = nn.LogSoftmax(dim=1)\n\n    def forward(self, inputs, targets):\n        \"\"\"\n        Args:\n            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)\n            targets: ground truth labels with shape (num_classes)\n        \"\"\"\n\n        log_probs = self.logsoftmax(inputs)\n        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).cpu(), 1)\n        if self.use_gpu: targets = targets.cuda()\n        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes\n        loss = (- targets * log_probs).mean(0).sum()\n        return loss\n\n\nfrom typing import Tuple\nfrom torch import Tensor\n\n\ndef convert_label_to_similarity(normed_feature: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:\n    similarity_matrix = normed_feature @ normed_feature.transpose(1, 0)\n    label_matrix = label.unsqueeze(1) == label.unsqueeze(0)\n\n    positive_matrix = label_matrix.triu(diagonal=1)\n    negative_matrix = label_matrix.logical_not().triu(diagonal=1)\n\n    similarity_matrix = similarity_matrix.view(-1)\n    positive_matrix = positive_matrix.view(-1)\n    negative_matrix = negative_matrix.view(-1)\n    return similarity_matrix[positive_matrix], similarity_matrix[negative_matrix]\n\n\ndef normalize(x, axis=-1):\n    x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)\n    return x\n\n\nclass CircleLoss(nn.Module):\n    def __init__(self, m: float, gamma: float) -> None:\n        super(CircleLoss, self).__init__()\n        self.m = m\n        self.gamma = gamma\n        self.soft_plus = nn.Softplus()\n\n    def forward(self, feat, label) -> Tensor:\n        feat = normalize(feat, axis=-1)\n        sp, sn = convert_label_to_similarity(feat, label)\n        ap = torch.clamp_min(- sp.detach() + 1 + self.m, min=0.)\n        an = torch.clamp_min(sn.detach() + self.m, min=0.)\n\n        delta_p = 1 - self.m\n        delta_n = self.m\n\n        logit_p = - ap * (sp - delta_p) * self.gamma\n        logit_n = an * (sn - delta_n) * self.gamma\n\n        loss = self.soft_plus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0))\n\n        return loss\n\n\nclass Checkpoint:\n    def __init__(self, ckpt):\n        self.ckpt = ckpt\n        self.init = -1000\n\n    def __call__(self, metric, model, optimizer):\n        if metric > self.init:\n            self.init = metric\n            torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, self.ckpt)\n\n\n\nclass EMA():\n    def __init__(self, model, decay):\n        self.model = model\n        self.decay = decay\n        self.shadow = {}\n        self.backup = {}\n\n    def register(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.data.clone()\n\n    def update(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                assert name in self.shadow\n                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]\n                self.shadow[name] = new_average.clone()\n\n    def apply_shadow(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                assert name in self.shadow\n                self.backup[name] = param.data\n                param.data = self.shadow[name]\n\n    def restore(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                assert name in self.backup\n                param.data = self.backup[name]\n        self.backup = {}\nfrom fastai.metrics import accuracy\ndef do_eval(model,valid_dl):\n    device=torch.device('cuda:0')\n    accs=0\n    with torch.no_grad():\n        for images,labels in valid_dl:\n            images=images.to(device)\n            labels=labels.to(device)\n            logit=model(images)\n            acc=accuracy(logit,labels)\n            accs+=acc\n    return accs/len(valid_dl)\n\n\n\nfrom torchtools.optim import RangerLars\nfrom torch.cuda.amp import autocast\nepochs=50\ntrain_dataset=Dataset(train_name,train_label,train_transform)\nval_dataset=Dataset(test_name,test_label,test_transform)\ntrain_dl=Data.DataLoader(train_dataset,batch_size=64,shuffle=True,num_workers=8)\nval_dl=Data.DataLoader(val_dataset,batch_size=64,shuffle=False,num_workers=8)\nmodel=Model('resnet50',10)\noptimizer=RangerLars(model.parameters(),lr=0.001)\nscheduler=Flat(optimizer,50)\ncircle_criterion=CircleLoss(m=0.25,gamma=80)\ncls_criterion=CrossEntropyLabelSmoothLoss(10)\n\nckpt='best.pt'\ncheckpoint=Checkpoint(ckpt)\nscaler=torch.cuda.amp.GradScaler()\ndevice=torch.device('cuda:0')\nmodel=model.to(device)\n\nfrom tqdm import tqdm\nwith tqdm(total=epochs) as pbar:\n    for epoch in range(epochs):\n        model.train()\n        for images,labels in train_dl:\n            \n            images=images.to(device)\n            labels=labels.to(device)\n            optimizer.zero_grad()\n            loss=0\n            with autocast():\n                feature,logit=model(images)\n                loss=loss+circle_criterion(feature,labels)\n                loss=loss+cls_criterion(logit,labels)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        scheduler.step()    \n            \n        if epoch%5==0:\n            model.eval()\n            acc=do_eval(model,val_dl)\n            print('epoch:{}--acc:{}'.format(epoch,acc))\n            checkpoint(acc,model,optimizer)\n        pbar.update(1)\n\n\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}