{"cells":[{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import os\nimport collections\nfrom datetime import datetime, timedelta\n\nos.environ[\"XRT_TPU_CONFIG\"] = \"tpu_worker;0;10.0.0.2:8470\"\n\n_VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')\nVERSION = \"torch_xla==nightly\"\nCONFIG = {\n    'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(\n        (datetime.today() - timedelta(1)).strftime('%Y%m%d')))}[VERSION]\n\nDIST_BUCKET = 'gs://tpu-pytorch/wheels'\nTORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\nTORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\nTORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n\n!export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH\n!apt-get install libomp5 -y\n!apt-get install libopenblas-dev -y\n\n!pip uninstall -y torch torchvision\n!gsutil cp \"$DIST_BUCKET/$TORCH_WHEEL\" .\n!gsutil cp \"$DIST_BUCKET/$TORCH_XLA_WHEEL\" .\n!gsutil cp \"$DIST_BUCKET/$TORCHVISION_WHEEL\" .\n!pip install \"$TORCH_WHEEL\"\n!pip install \"$TORCH_XLA_WHEEL\"\n!pip install \"$TORCHVISION_WHEEL\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport re\nimport cv2\nimport time\nimport tensorflow\nimport collections\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom glob import glob\nfrom PIL import Image\nimport requests, threading\nimport matplotlib.pyplot as plt\nfrom datetime import datetime, timedelta","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torchvision import datasets\nfrom torchvision import transforms\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import OneCycleLR","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\ntorch.manual_seed(42)\ntorch.set_default_tensor_type('torch.FloatTensor')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"DATASET_DIR = '/kaggle/input/104-flowers-garden-of-eden/jpeg-224x224/'\nTRAIN_DIR  = DATASET_DIR + '/train'\nVAL_DIR  = DATASET_DIR + '/val'\nTEST_DIR  = DATASET_DIR + '/test'\n\nBATCH_SIZE = 128 # per core\nNUM_EPOCH = 3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n\ntrain_transform = transforms.Compose([transforms.RandomResizedCrop(224),\n                                      transforms.RandomHorizontalFlip(0.5),\n                                      transforms.ToTensor(),\n                                      normalize])\n\nvalid_transform = transforms.Compose([transforms.Resize((224,224)),\n                                      transforms.ToTensor(),\n                                      normalize])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = datasets.ImageFolder(TRAIN_DIR, transform=train_transform)\nvalid = datasets.ImageFolder(VAL_DIR, transform=valid_transform)\n\n# print out some data stats\nprint('Num training images: ', len(train))\nprint('Num test images: ', len(valid))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n    \"\"\"3x3 convolution with padding\"\"\"\n    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n                     padding=dilation, groups=groups, bias=False, dilation=dilation)\n\n\ndef conv1x1(in_planes, out_planes, stride=1):\n    \"\"\"1x1 convolution\"\"\"\n    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n    __constants__ = ['downsample']\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n                 base_width=64, dilation=1, norm_layer=None):\n        super(BasicBlock, self).__init__()\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        if groups != 1 or base_width != 64:\n            raise ValueError('BasicBlock only supports groups=1 and base_width=64')\n        if dilation > 1:\n            raise NotImplementedError(\"Dilation > 1 not supported in BasicBlock\")\n        # Both self.conv1 and self.downsample layers downsample the input when stride != 1\n        self.conv1 = conv3x3(inplanes, planes, stride)\n        self.bn1 = norm_layer(planes)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3(planes, planes)\n        self.bn2 = norm_layer(planes)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        identity = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        if self.downsample is not None:\n            identity = self.downsample(x)\n\n        out += identity\n        out = self.relu(out)\n\n        return out\n\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n    __constants__ = ['downsample']\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n                 base_width=64, dilation=1, norm_layer=None):\n        super(Bottleneck, self).__init__()\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        width = int(planes * (base_width / 64.)) * groups\n        # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n        self.conv1 = conv1x1(inplanes, width)\n        self.bn1 = norm_layer(width)\n        self.conv2 = conv3x3(width, width, stride, groups, dilation)\n        self.bn2 = norm_layer(width)\n        self.conv3 = conv1x1(width, planes * self.expansion)\n        self.bn3 = norm_layer(planes * self.expansion)\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        identity = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        if self.downsample is not None:\n            identity = self.downsample(x)\n\n        out += identity\n        out = self.relu(out)\n\n        return out\n\n\nclass ResNet(nn.Module):\n\n    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,\n                 groups=1, width_per_group=64, replace_stride_with_dilation=None,\n                 norm_layer=None):\n        super(ResNet, self).__init__()\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        self._norm_layer = norm_layer\n\n        self.inplanes = 64\n        self.dilation = 1\n        if replace_stride_with_dilation is None:\n            # each element in the tuple indicates if we should replace\n            # the 2x2 stride with a dilated convolution instead\n            replace_stride_with_dilation = [False, False, False]\n        if len(replace_stride_with_dilation) != 3:\n            raise ValueError(\"replace_stride_with_dilation should be None \"\n                             \"or a 3-element tuple, got {}\".format(replace_stride_with_dilation))\n        self.groups = groups\n        self.base_width = width_per_group\n        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,\n                               bias=False)\n        self.bn1 = norm_layer(self.inplanes)\n        self.relu = nn.ReLU(inplace=True)\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n        self.layer1 = self._make_layer(block, 64, layers[0])\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,\n                                       dilate=replace_stride_with_dilation[0])\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,\n                                       dilate=replace_stride_with_dilation[1])\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,\n                                       dilate=replace_stride_with_dilation[2])\n        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n        self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n        # Zero-initialize the last BN in each residual branch,\n        # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n        if zero_init_residual:\n            for m in self.modules():\n                if isinstance(m, Bottleneck):\n                    nn.init.constant_(m.bn3.weight, 0)\n                elif isinstance(m, BasicBlock):\n                    nn.init.constant_(m.bn2.weight, 0)\n\n    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):\n        norm_layer = self._norm_layer\n        downsample = None\n        previous_dilation = self.dilation\n        if dilate:\n            self.dilation *= stride\n            stride = 1\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                conv1x1(self.inplanes, planes * block.expansion, stride),\n                norm_layer(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,\n                            self.base_width, previous_dilation, norm_layer))\n        self.inplanes = planes * block.expansion\n        for _ in range(1, blocks):\n            layers.append(block(self.inplanes, planes, groups=self.groups,\n                                base_width=self.base_width, dilation=self.dilation,\n                                norm_layer=norm_layer))\n\n        return nn.Sequential(*layers)\n\n    def _forward_impl(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        x = self.avgpool(x)\n        x = torch.flatten(x, 1)\n        x = self.fc(x)\n\n        return x\n\n    def forward(self, x):\n        return self._forward_impl(x)\n\n\ndef _resnet(arch, block, layers, pretrained, progress, **kwargs):\n    model = ResNet(block, layers, **kwargs)\n    if pretrained:\n        state_dict = load_state_dict_from_url(model_urls[arch],\n                                              progress=progress)\n        model.load_state_dict(state_dict)\n    return model\n\n\ndef resnet34(pretrained=False, progress=True, **kwargs):\n    r\"\"\"ResNet-34 model from\n    `\"Deep Residual Learning for Image Recognition\" <https://arxiv.org/pdf/1512.03385.pdf>`_\n    Args:\n        pretrained (bool): If True, returns a model pre-trained on ImageNet\n        progress (bool): If True, displays a progress bar of the download to stderr\n    \"\"\"\n    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,\n                   **kwargs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_model():\n    global train, valid\n    \n    torch.manual_seed(42)\n    \n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True)\n    \n    train_loader = torch.utils.data.DataLoader(\n        train,\n        batch_size=BATCH_SIZE,\n        sampler=train_sampler,\n        num_workers=0,\n        drop_last=True) # print(len(train_loader))\n    \n    valid_loader = torch.utils.data.DataLoader(\n        valid,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=0,\n        drop_last=True)\n    \n    xm.master_print(f\"Train for {len(train_loader)} steps per epoch\")\n    # Scale learning rate to num cores\n    learning_rate = 0.0001 * xm.xrt_world_size()\n\n    # Get loss function, optimizer, and model\n    device = xm.xla_device()\n    model = resnet34(num_classes=104, pretrained=False).to(device)\n    \n    model = model.to(device)\n    loss_fn =  nn.CrossEntropyLoss()\n    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=5e-4)\n    scheduler = OneCycleLR(optimizer, \n                           learning_rate, \n                           div_factor=10.0, \n                           final_div_factor=50.0, \n                           epochs=NUM_EPOCH,\n                           steps_per_epoch=len(train_loader))\n    \n    \n    \n    def train_loop_fn(loader):\n        tracker = xm.RateTracker()\n        model.train()\n        for x, (data, target) in enumerate(loader):\n            optimizer.zero_grad()\n            output = model(data)\n            loss = loss_fn(output, target)\n            loss.backward()\n            xm.optimizer_step(optimizer)\n            tracker.add(data.shape[0])\n            scheduler.step()\n            if x % 30 == 0:\n                print('[xla:{}]({})\\tLoss={:.3f}\\tRate={:.2f}\\tGlobalRate={:.2f}'.format(\n                    xm.get_ordinal(), x, loss.item(), tracker.rate(),\n                    tracker.global_rate()), flush=True)\n\n    def test_loop_fn(loader):\n        with torch.no_grad():\n            total_samples, correct = 0, 0\n            model.eval()\n            for data, target in loader:\n                output = model(data)\n                pred = output.max(1, keepdim=True)[1]\n                correct += pred.eq(target.view_as(pred)).sum().item()\n                total_samples += data.size()[0]\n            accuracy = 100.0 * correct / total_samples\n            print('[xla:{}] Accuracy={:.2f}%'.format(xm.get_ordinal(), accuracy), flush=True)\n            model.train()\n        return accuracy\n\n    # Train and eval loops\n    accuracy = []\n    for epoch in range(1, NUM_EPOCH + 1):\n        start = time.time()\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        train_loop_fn(para_loader.per_device_loader(device))\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        accuracy.append(test_loop_fn(para_loader.per_device_loader(device)))\n        xm.save(model.state_dict(), \"./model.pt\")\n        xm.master_print(\"Finished training epoch {} acc {:.2f} in {:.2f} sec\\n\"\\\n                        .format(epoch, accuracy[-1], time.time() - start))        \n        \n        if epoch == 15: #unfreeze\n                for param in model.base_model.parameters():\n                    param.requires_grad = True\n\n    return accuracy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    global acc_list\n    torch.set_default_tensor_type('torch.FloatTensor')\n    res = train_model()\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.6"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"093bddf9af6248a68e174d739cbe40e8":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":"initial"}},"79838c22264849ab88666a071b439caf":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_de7e4ee4e19e4b548575d5527d6bb66c","IPY_MODEL_91a6f5abd3d046649428393bde8bec19"],"layout":"IPY_MODEL_8af46dff9071427088dbe6088dc5ab5a"}},"79db4cb5de884e6ab48c171b3bfc8667":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"8a880d0a57034166b3644fbda4d301fa":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"8af46dff9071427088dbe6088dc5ab5a":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"91a6f5abd3d046649428393bde8bec19":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_8a880d0a57034166b3644fbda4d301fa","placeholder":"​","style":"IPY_MODEL_a7108d767f13414bb2770815731423bc","value":" 77.4M/77.4M [00:00&lt;00:00, 92.8MB/s]"}},"a7108d767f13414bb2770815731423bc":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"de7e4ee4e19e4b548575d5527d6bb66c":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"100%","description_tooltip":null,"layout":"IPY_MODEL_79db4cb5de884e6ab48c171b3bfc8667","max":81131730,"min":0,"orientation":"horizontal","style":"IPY_MODEL_093bddf9af6248a68e174d739cbe40e8","value":81131730}}},"version_major":2,"version_minor":0}}},"nbformat":4,"nbformat_minor":4}