{"cells":[{"metadata":{},"cell_type":"markdown","source":"# TPU setup"},{"metadata":{"trusted":true},"cell_type":"code","source":"VERSION = \"nightly\"  #@param [\"1.5\" , \"20200325\", \"nightly\"]\n!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version $VERSION --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://www.kaggle.com/c/flower-classification-with-tpus/discussion/129820\nimport os\nos.environ['XRT_TPU_CONFIG'] = \"tpu_worker;0;10.0.0.2:8470\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Library"},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np \nimport pandas as pd \nimport json","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"os.listdir('../input/104-flowers-garden-of-eden/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.DataFrame()\nIMG_FOLDER = '512x512'\n\n# train/\nimage_dict = {}\nfor _class in os.listdir(f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/train/'):\n    image_dict[_class] = os.listdir(f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/train/{_class}/')   \nfor k, values in image_dict.items():\n    train = pd.concat([train, pd.DataFrame({'id': [f'train/{k}/{v}' for v in values], 'class': [k]*len(values)})])\n\n# val/\nimage_dict = {}\nfor _class in os.listdir(f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/val/'):\n    image_dict[_class] = os.listdir(f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/val/{_class}/')   \nfor k, values in image_dict.items():\n    train = pd.concat([train, pd.DataFrame({'id': [f'val/{k}/{v}' for v in values], 'class': [k]*len(values)})])\n\ntrain = train.reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# From https://www.kaggle.com/mgornergoogle/getting-started-with-100-flowers-on-tpu\nCLASSES = ['pink primrose',    'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea',     'wild geranium',     'tiger lily',           'moon orchid',              'bird of paradise', 'monkshood',        'globe thistle',         # 00 - 09\n           'snapdragon',       \"colt's foot\",               'king protea',      'spear thistle', 'yellow iris',       'globe-flower',         'purple coneflower',        'peruvian lily',    'balloon flower',   'giant white arum lily', # 10 - 19\n           'fire lily',        'pincushion flower',         'fritillary',       'red ginger',    'grape hyacinth',    'corn poppy',           'prince of wales feathers', 'stemless gentian', 'artichoke',        'sweet william',         # 20 - 29\n           'carnation',        'garden phlox',              'love in the mist', 'cosmos',        'alpine sea holly',  'ruby-lipped cattleya', 'cape flower',              'great masterwort', 'siam tulip',       'lenten rose',           # 30 - 39\n           'barberton daisy',  'daffodil',                  'sword lily',       'poinsettia',    'bolero deep blue',  'wallflower',           'marigold',                 'buttercup',        'daisy',            'common dandelion',      # 40 - 49\n           'petunia',          'wild pansy',                'primula',          'sunflower',     'lilac hibiscus',    'bishop of llandaff',   'gaura',                    'geranium',         'orange dahlia',    'pink-yellow dahlia',    # 50 - 59\n           'cautleya spicata', 'japanese anemone',          'black-eyed susan', 'silverbush',    'californian poppy', 'osteospermum',         'spring crocus',            'iris',             'windflower',       'tree poppy',            # 60 - 69\n           'gazania',          'azalea',                    'water lily',       'rose',          'thorn apple',       'morning glory',        'passion flower',           'lotus',            'toad lily',        'anthurium',             # 70 - 79\n           'frangipani',       'clematis',                  'hibiscus',         'columbine',     'desert-rose',       'tree mallow',          'magnolia',                 'cyclamen ',        'watercress',       'canna lily',            # 80 - 89\n           'hippeastrum ',     'bee balm',                  'pink quill',       'foxglove',      'bougainvillea',     'camellia',             'mallow',                   'mexican petunia',  'bromelia',         'blanket flower',        # 90 - 99\n           'trumpet creeper',  'blackberry lily',           'common tulip',     'wild rose']                                                                                                                                               # 100 - 102","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class_map = {}\n\nfor i, c in enumerate(CLASSES):\n    class_map[c] = i\n    \ntrain['class'] = train['class'].map(class_map)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Loading"},{"metadata":{"trusted":true},"cell_type":"code","source":"#submission = pd.read_csv('../input/flower-classification-with-tpus/sample_submission.csv')\nsubmission = pd.read_csv('../input/getting-started-with-100-flowers-on-tpu/submission.csv').drop(columns='label')\nsubmission.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Library"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Library\n# ====================================================\n\nimport sys\n\nimport gc\nimport os\nimport random\nimport time\nfrom contextlib import contextmanager\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\n\nimport cv2\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nimport scipy as sp\n\nimport sklearn.metrics\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold\n\nfrom functools import partial\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam, SGD\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.models as models\n\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, IAAAdditiveGaussianNoise\n)\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\n\n\n#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n#device","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla\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.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\n\n@contextmanager\ndef timer(name):\n    t0 = time.time()\n    LOGGER.info(f'[{name}] start')\n    yield\n    LOGGER.info(f'[{name}] done in {time.time() - t0:.0f} s.')\n\n    \ndef init_logger(log_file='train.log'):\n    from logging import getLogger, DEBUG, FileHandler,  Formatter,  StreamHandler\n    \n    log_format = '%(asctime)s %(levelname)s %(message)s'\n    \n    stream_handler = StreamHandler()\n    stream_handler.setLevel(DEBUG)\n    stream_handler.setFormatter(Formatter(log_format))\n    \n    file_handler = FileHandler(log_file)\n    file_handler.setFormatter(Formatter(log_format))\n    \n    logger = getLogger('Flower')\n    logger.setLevel(DEBUG)\n    logger.addHandler(stream_handler)\n    logger.addHandler(file_handler)\n    \n    return logger\n\nLOG_FILE = 'train.log'\nLOGGER = init_logger(LOG_FILE)\n\n\ndef seed_torch(seed=777):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nSEED = 777\nseed_torch(SEED)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"N_CLASSES = 104\n\n\nclass TrainDataset(Dataset):\n    def __init__(self, df, labels, transform=None):\n        self.df = df\n        self.labels = labels\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.df['id'].values[idx]\n        file_path = f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/{file_name}'\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n            \n        label = self.labels.values[idx]\n        \n        return image, label\n    \n    \nclass TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.df['id'].values[idx]\n        file_path = f'../input/104-flowers-garden-of-eden/jpeg-{IMG_FOLDER}/test/{file_name}.jpeg'\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        \n        return image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Transforms"},{"metadata":{"trusted":true},"cell_type":"code","source":"HEIGHT = 256\nWIDTH = 256\n\n\ndef get_transforms(*, data):\n    \n    assert data in ('train', 'valid', 'test-tta')\n    \n    if data == 'train':\n        return Compose([\n            #Resize(HEIGHT, WIDTH),\n            RandomResizedCrop(HEIGHT, WIDTH),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(rotate_limit=30, p=0.5),\n            Cutout(p=0.5, max_h_size=12, max_w_size=12, num_holes=6),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    \n    elif data == 'valid':\n        return Compose([\n            Resize(HEIGHT, WIDTH),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    \n    elif data == 'test-tta':\n        return Compose([\n            RandomResizedCrop(HEIGHT, WIDTH),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# train valid split"},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.model_selection import StratifiedKFold\n\nDEBUG = False\nN_FOLD = 4\n\nif DEBUG:\n    folds = train.sample(n=1000, random_state=42).reset_index(drop=True).copy()\nelse:\n    folds = train.copy()\n    \ntrain_labels = folds['class'].values\nkf = StratifiedKFold(n_splits=N_FOLD, shuffle=True, random_state=42)\nfor fold, (train_index, val_index) in enumerate(kf.split(folds.values, train_labels)):\n    folds.loc[val_index, 'fold'] = int(fold)\nfolds['fold'] = folds['fold'].astype(int)\n\nfolds.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/Cadene/pretrained-models.pytorch/blob/master/pretrainedmodels/models/senet.py\n\nfrom collections import OrderedDict\nimport math\n\n\nclass SEModule(nn.Module):\n\n    def __init__(self, channels, reduction):\n        super(SEModule, self).__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1,\n                             padding=0)\n        self.relu = nn.ReLU(inplace=True)\n        self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1,\n                             padding=0)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        module_input = x\n        x = self.avg_pool(x)\n        x = self.fc1(x)\n        x = self.relu(x)\n        x = self.fc2(x)\n        x = self.sigmoid(x)\n        return module_input * x\n\n\nclass Bottleneck(nn.Module):\n    \"\"\"\n    Base class for bottlenecks that implements `forward()` method.\n    \"\"\"\n    def forward(self, x):\n        residual = 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            residual = self.downsample(x)\n\n        out = self.se_module(out) + residual\n        out = self.relu(out)\n\n        return out\n\n\nclass SEBottleneck(Bottleneck):\n    \"\"\"\n    Bottleneck for SENet154.\n    \"\"\"\n    expansion = 4\n\n    def __init__(self, inplanes, planes, groups, reduction, stride=1,\n                 downsample=None):\n        super(SEBottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes * 2)\n        self.conv2 = nn.Conv2d(planes * 2, planes * 4, kernel_size=3,\n                               stride=stride, padding=1, groups=groups,\n                               bias=False)\n        self.bn2 = nn.BatchNorm2d(planes * 4)\n        self.conv3 = nn.Conv2d(planes * 4, planes * 4, kernel_size=1,\n                               bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n        self.relu = nn.ReLU(inplace=True)\n        self.se_module = SEModule(planes * 4, reduction=reduction)\n        self.downsample = downsample\n        self.stride = stride\n\n\nclass SEResNetBottleneck(Bottleneck):\n    \"\"\"\n    ResNet bottleneck with a Squeeze-and-Excitation module. It follows Caffe\n    implementation and uses `stride=stride` in `conv1` and not in `conv2`\n    (the latter is used in the torchvision implementation of ResNet).\n    \"\"\"\n    expansion = 4\n\n    def __init__(self, inplanes, planes, groups, reduction, stride=1,\n                 downsample=None):\n        super(SEResNetBottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False,\n                               stride=stride)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1,\n                               groups=groups, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n        self.relu = nn.ReLU(inplace=True)\n        self.se_module = SEModule(planes * 4, reduction=reduction)\n        self.downsample = downsample\n        self.stride = stride\n\n\nclass SEResNeXtBottleneck(Bottleneck):\n    \"\"\"\n    ResNeXt bottleneck type C with a Squeeze-and-Excitation module.\n    \"\"\"\n    expansion = 4\n\n    def __init__(self, inplanes, planes, groups, reduction, stride=1,\n                 downsample=None, base_width=4):\n        super(SEResNeXtBottleneck, self).__init__()\n        width = math.floor(planes * (base_width / 64)) * groups\n        self.conv1 = nn.Conv2d(inplanes, width, kernel_size=1, bias=False,\n                               stride=1)\n        self.bn1 = nn.BatchNorm2d(width)\n        self.conv2 = nn.Conv2d(width, width, kernel_size=3, stride=stride,\n                               padding=1, groups=groups, bias=False)\n        self.bn2 = nn.BatchNorm2d(width)\n        self.conv3 = nn.Conv2d(width, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n        self.relu = nn.ReLU(inplace=True)\n        self.se_module = SEModule(planes * 4, reduction=reduction)\n        self.downsample = downsample\n        self.stride = stride\n\n\nclass SENet(nn.Module):\n\n    def __init__(self, block, layers, groups, reduction, dropout_p=0.2,\n                 inplanes=128, input_3x3=True, downsample_kernel_size=3,\n                 downsample_padding=1, num_classes=1000):\n        super(SENet, self).__init__()\n        self.inplanes = inplanes\n        if input_3x3:\n            layer0_modules = [\n                ('conv1', nn.Conv2d(3, 64, 3, stride=2, padding=1,\n                                    bias=False)),\n                ('bn1', nn.BatchNorm2d(64)),\n                ('relu1', nn.ReLU(inplace=True)),\n                ('conv2', nn.Conv2d(64, 64, 3, stride=1, padding=1,\n                                    bias=False)),\n                ('bn2', nn.BatchNorm2d(64)),\n                ('relu2', nn.ReLU(inplace=True)),\n                ('conv3', nn.Conv2d(64, inplanes, 3, stride=1, padding=1,\n                                    bias=False)),\n                ('bn3', nn.BatchNorm2d(inplanes)),\n                ('relu3', nn.ReLU(inplace=True)),\n            ]\n        else:\n            layer0_modules = [\n                ('conv1', nn.Conv2d(3, inplanes, kernel_size=7, stride=2,\n                                    padding=3, bias=False)),\n                ('bn1', nn.BatchNorm2d(inplanes)),\n                ('relu1', nn.ReLU(inplace=True)),\n            ]\n        # To preserve compatibility with Caffe weights `ceil_mode=True`\n        # is used instead of `padding=1`.\n        layer0_modules.append(('pool', nn.MaxPool2d(3, stride=2,\n                                                    ceil_mode=True)))\n        self.layer0 = nn.Sequential(OrderedDict(layer0_modules))\n        self.layer1 = self._make_layer(\n            block,\n            planes=64,\n            blocks=layers[0],\n            groups=groups,\n            reduction=reduction,\n            downsample_kernel_size=1,\n            downsample_padding=0\n        )\n        self.layer2 = self._make_layer(\n            block,\n            planes=128,\n            blocks=layers[1],\n            stride=2,\n            groups=groups,\n            reduction=reduction,\n            downsample_kernel_size=downsample_kernel_size,\n            downsample_padding=downsample_padding\n        )\n        self.layer3 = self._make_layer(\n            block,\n            planes=256,\n            blocks=layers[2],\n            stride=2,\n            groups=groups,\n            reduction=reduction,\n            downsample_kernel_size=downsample_kernel_size,\n            downsample_padding=downsample_padding\n        )\n        self.layer4 = self._make_layer(\n            block,\n            planes=512,\n            blocks=layers[3],\n            stride=2,\n            groups=groups,\n            reduction=reduction,\n            downsample_kernel_size=downsample_kernel_size,\n            downsample_padding=downsample_padding\n        )\n        self.avg_pool = nn.AvgPool2d(7, stride=1)\n        self.dropout = nn.Dropout(dropout_p) if dropout_p is not None else None\n        self.last_linear = nn.Linear(512 * block.expansion, num_classes)\n\n    def _make_layer(self, block, planes, blocks, groups, reduction, stride=1,\n                    downsample_kernel_size=1, downsample_padding=0):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.inplanes, planes * block.expansion,\n                          kernel_size=downsample_kernel_size, stride=stride,\n                          padding=downsample_padding, bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, groups, reduction, stride,\n                            downsample))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes, groups, reduction))\n\n        return nn.Sequential(*layers)\n\n    def features(self, x):\n        x = self.layer0(x)\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n        return x\n\n    def logits(self, x):\n        x = self.avg_pool(x)\n        if self.dropout is not None:\n            x = self.dropout(x)\n        x = x.view(x.size(0), -1)\n        x = self.last_linear(x)\n        return x\n\n    def forward(self, x):\n        x = self.features(x)\n        x = self.logits(x)\n        return x\n\n\ndef initialize_pretrained_model(model, num_classes, settings):\n    assert num_classes == settings['num_classes'], \\\n        'num_classes should be {}, but is {}'.format(\n            settings['num_classes'], num_classes)\n    model.load_state_dict(model_zoo.load_url(settings['url']))\n    model.input_space = settings['input_space']\n    model.input_size = settings['input_size']\n    model.input_range = settings['input_range']\n    model.mean = settings['mean']\n    model.std = settings['std']\n\n\ndef se_resnext50_32x4d(num_classes=1000, pretrained='imagenet'):\n    model = SENet(SEResNeXtBottleneck, [3, 4, 6, 3], groups=32, reduction=16,\n                  dropout_p=None, inplanes=64, input_3x3=False,\n                  downsample_kernel_size=1, downsample_padding=0,\n                  num_classes=num_classes)\n    if pretrained is not None:\n        settings = pretrained_settings['se_resnext50_32x4d'][pretrained]\n        initialize_pretrained_model(model, num_classes, settings)\n    return model\n\n\ndef se_resnext101_32x4d(num_classes=1000, pretrained='imagenet'):\n    model = SENet(SEResNeXtBottleneck, [3, 4, 23, 3], groups=32, reduction=16,\n                  dropout_p=None, inplanes=64, input_3x3=False,\n                  downsample_kernel_size=1, downsample_padding=0,\n                  num_classes=num_classes)\n    if pretrained is not None:\n        settings = pretrained_settings['se_resnext101_32x4d'][pretrained]\n        initialize_pretrained_model(model, num_classes, settings)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pretrained_path = {'se_resnext50_32x4d': '../input/pytorch-se-resnext/se_resnext50_32x4d-a260b3a4.pth',\n                   'se_resnext101_32x4d': '../input/pytorch-se-resnext/se_resnext101_32x4d-3b2fe3d8.pth',}\n\nclass CustomSEResNeXt(nn.Module):\n\n    def __init__(self, model_name='se_resnext50_32x4d'):\n        assert model_name in ('se_resnext50_32x4d', 'se_resnext101_32x4d')\n        super().__init__()\n        \n        self.model = se_resnext50_32x4d(pretrained=None)\n        weights_path = pretrained_path[model_name]\n        self.model.load_state_dict(torch.load(weights_path))\n        self.model.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.model.last_linear = nn.Linear(self.model.last_linear.in_features, N_CLASSES)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import f1_score","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n    \n        \n    def train_loop_fn(para_train_loader, model, optimizer, criterion, device):\n            \n        model.train()\n        avg_loss = 0.\n\n        optimizer.zero_grad()\n            \n        for i, (images, labels) in enumerate(para_train_loader.per_device_loader(device)):\n\n            images = images.to(device)\n            labels = labels.to(device)\n\n            y_preds = model(images)\n            loss = criterion(y_preds, labels)\n            \n            if i % 40 == 0:\n                xm.master_print(f'[train] i={i}, loss={loss}')\n                    \n            loss.backward()\n            #optimizer.step()\n            xm.optimizer_step(optimizer)\n            optimizer.zero_grad()\n\n            avg_loss += loss.item() / len(train_loader)\n                \n        return avg_loss\n          \n        \n    def eval_loop_fn(para_valid_loader, model, criterion, device, scheduler):\n\n        model.eval()\n        avg_val_loss = 0.\n        preds = []\n        valid_labels = []\n            \n        for i, (images, labels) in enumerate(para_valid_loader.per_device_loader(device)):\n\n            images = images.to(device)\n            labels = labels.to(device)\n\n            with torch.no_grad():\n                y_preds = model(images)\n\n            preds.append(y_preds.argmax(1).to('cpu').numpy())\n            valid_labels.append(labels.to('cpu').numpy())\n\n            loss = criterion(y_preds, labels)\n            avg_val_loss += loss.item() / len(valid_loader)\n\n        scheduler.step(avg_val_loss)\n\n        preds = np.concatenate(preds)\n        valid_labels = np.concatenate(valid_labels)\n        \n        score = f1_score(valid_labels, preds, average='macro')\n            \n        return avg_val_loss, score\n        \n        \n    device = xm.xla_device()\n    world_size = xm.xrt_world_size()\n        \n    batch_size = int( 128 / world_size )\n    n_epochs = 20\n    lr = 1e-4 * world_size\n        \n    xm.master_print(f'device: {device}')\n    xm.master_print(f'world_size: {world_size}')\n    xm.master_print(f'batch_size: {batch_size}')\n    xm.master_print(f'n_epochs: {n_epochs}')\n    xm.master_print(f'lr: {lr}')\n    \n    NUM_TTA = 1\n    ENSEMBLE_WEIGHTS = {'se_resnext': 1}\n    probas = []\n    model1_proba = []\n    \n    for FOLD in range(N_FOLD):\n        \n        xm.master_print(f\"FOLD: {FOLD}\")\n        \n        trn_idx = folds[folds['fold'] != FOLD].index\n        val_idx = folds[folds['fold'] == FOLD].index\n        \n        train_dataset = TrainDataset(folds.loc[trn_idx].reset_index(drop=True), \n                             folds.loc[trn_idx]['class'], \n                             transform=get_transforms(data='train'))\n        valid_dataset = TrainDataset(folds.loc[val_idx].reset_index(drop=True), \n                                     folds.loc[val_idx]['class'], \n                                     transform=get_transforms(data='valid'))\n        test_dataset = TestDataset(submission, transform=get_transforms(data='valid'))\n\n        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,\n                                                                        num_replicas=xm.xrt_world_size(),\n                                                                        rank=xm.get_ordinal(),\n                                                                        shuffle=True)\n        valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_dataset,\n                                                                        num_replicas=xm.xrt_world_size(),\n                                                                        rank=xm.get_ordinal(),\n                                                                        shuffle=False)\n        test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset,\n                                                                       num_replicas=xm.xrt_world_size(),\n                                                                       rank=xm.get_ordinal(),\n                                                                       shuffle=False)\n        if world_size==1:\n            N_JOBS = 4\n        else:\n            N_JOBS = 0\n        train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, drop_last=True, num_workers=N_JOBS)\n        valid_loader = DataLoader(valid_dataset, batch_size=batch_size, sampler=valid_sampler, drop_last=False, num_workers=N_JOBS)\n        test_loader = DataLoader(test_dataset, batch_size=batch_size, sampler=test_sampler, drop_last=False, num_workers=N_JOBS)\n        \n        xm.master_print(f\"Train for {len(train_loader)} steps per epoch\")\n        \n        my_model = CustomSEResNeXt(model_name='se_resnext50_32x4d')\n        model = my_model.to(device)\n\n        optimizer = Adam(model.parameters(), lr=lr, amsgrad=False)\n        scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.75, patience=1, verbose=True, eps=1e-6)\n\n        criterion = nn.CrossEntropyLoss()\n        best_loss = np.inf\n        best_score = 0.\n        best_thresh = 0.\n\n        for epoch in range(n_epochs):\n\n            start_time = time.time()\n            \n            # train\n            para_train_loader = pl.ParallelLoader(train_loader, [device])\n            avg_loss = train_loop_fn(para_train_loader, model, optimizer, criterion, device)\n            \n            # eval\n            para_valid_loader = pl.ParallelLoader(valid_loader, [device])\n            avg_val_loss, score = eval_loop_fn(para_valid_loader, model, criterion, device, scheduler)\n\n            elapsed = time.time() - start_time\n\n            if world_size==1:\n                LOGGER.debug(f'  Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n                LOGGER.debug(f'  Epoch {epoch+1} - f1_score: {score}')\n            else:\n                xm.master_print(f'  Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n                xm.master_print(f'  Epoch {epoch+1} - f1_score: {score}')\n\n            if score>best_score:\n                best_score = score\n                if world_size==1:\n                    LOGGER.debug(f'  Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n                    torch.save({'model': model.state_dict(), \n                                'score': best_score}, \n                               f'se_resnext_fold{FOLD}.pth')\n                else:\n                    xm.master_print(f'  Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n                    xm.save({'model': model.state_dict(), \n                             'score': best_score}, \n                            f'se_resnext_fold{FOLD}.pth')\n                    \n        # inference\n        model.load_state_dict(torch.load(f'./se_resnext_fold{FOLD}.pth')['model'])\n\n        for param in model.parameters():\n            param.requires_grad = False\n        model.eval()\n\n        TTA_list = []\n\n        for _ in range(NUM_TTA):\n\n            test_proba = []\n\n            for i, images in enumerate(test_loader):\n\n                images = images.to(device) \n\n                with torch.no_grad():\n                    y_preds = model(images)\n\n                test_proba.append(list(y_preds.to('cpu').numpy()))\n                \n                if i % 20 == 0:\n                    print(f'[inference] i={i}... done')\n\n            TTA_list.append(sum(test_proba, []))\n\n        # TTA average\n        model1_proba.append(np.mean([test_preds for test_preds in TTA_list], axis=0))\n\n    # FOLD average\n    probas.append(np.mean(model1_proba, axis=0))\n    \n    ensemble_proba = ENSEMBLE_WEIGHTS['se_resnext']*np.array(probas[0])\n        \n    # ensemble predictions\n    predictions = ensemble_proba.argmax(1)\n    \n    submission['label'] = predictions\n    submission['label'] = submission['label'].astype(int)\n    submission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = _run()\n\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, 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":{"3152668ea8b74589b2a7123dccca5cbd":{"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"}},"6fd44a996a19464f99110602a6db08b7":{"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}},"79a7f6b8f5914a688c09d843603e8bc8":{"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_b66bcff488cd470c9bb25673405f2d82","max":46827520,"min":0,"orientation":"horizontal","style":"IPY_MODEL_3152668ea8b74589b2a7123dccca5cbd","value":46827520}},"83e67fce80104b5a9383ec9c3196506a":{"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}},"b66bcff488cd470c9bb25673405f2d82":{"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}},"d5e26b2104c94d3e93c1579f9830e0c5":{"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_79a7f6b8f5914a688c09d843603e8bc8","IPY_MODEL_f52096278de8485f9182c22905489f44"],"layout":"IPY_MODEL_83e67fce80104b5a9383ec9c3196506a"}},"f52096278de8485f9182c22905489f44":{"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_6fd44a996a19464f99110602a6db08b7","placeholder":"​","style":"IPY_MODEL_fdc2f0a4a2ad4377affda7d6294960c7","value":" 44.7M/44.7M [00:00<00:00, 70.6MB/s]"}},"fdc2f0a4a2ad4377affda7d6294960c7":{"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":""}}},"version_major":2,"version_minor":0}}},"nbformat":4,"nbformat_minor":4}