{"cells":[{"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":"!pip install pretrainedmodels\n!pip install efficientnet_pytorch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimport json\nimport sys\nimport gc\nimport os\nimport random\nimport time\nimport cv2\n\nfrom contextlib import contextmanager\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\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, f1_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.utils import model_zoo\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\nimport 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\n\nfrom collections import OrderedDict\nimport math\n\nimport pretrainedmodels as pmodels\nfrom efficientnet_pytorch import EfficientNet\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\nclass FlowerDataset(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    \ndef get_transforms(*, data):\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\nos.environ['XRT_TPU_CONFIG'] = \"tpu_worker;0;10.0.0.2:8470\"\n\nos.listdir('../input/104-flowers-garden-of-eden/')\n\nSEED = 777\nseed_torch(SEED)\n\nN_CLASSES = 104\n\nHEIGHT = 256\nWIDTH = 256\n\ntrain = pd.DataFrame()\nvalid = 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    valid = pd.concat([valid, pd.DataFrame({'id': [f'val/{k}/{v}' for v in values], 'class': [k]*len(values)})])\n\ntrain = train.reset_index(drop=True)\nvalid = valid.reset_index(drop=True)\n\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\n\nclass_map = {}\n\nfor i, c in enumerate(CLASSES):\n    class_map[c] = i\n    \ntrain['class'] = train['class'].map(class_map)\nvalid['class'] = valid['class'].map(class_map)\ntrain.head()\nvalid.head()\n\nDEBUG = False\n\nif DEBUG:\n    trn_folds = train.sample(n=1000, random_state=42).reset_index(drop=True).copy()\n    val_folds = valid.sample(n=1000, random_state=42).reset_index(drop=True).copy()\nelse:\n    trn_folds = train.copy()\n    val_folds = valid.copy()\n\ntrn_folds.head()\nval_folds.head()\n\ndef _run():\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            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    \n    trn_idx = trn_folds.index\n    val_idx = val_folds.index\n        \n    train_dataset = FlowerDataset(trn_folds.loc[trn_idx].reset_index(drop=True), \n                             trn_folds.loc[trn_idx]['class'], \n                             transform=get_transforms(data='train'))\n        \n    valid_dataset = FlowerDataset(val_folds.loc[val_idx].reset_index(drop=True), \n                                     val_folds.loc[val_idx]['class'],\n                                     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    if world_size==1:\n        N_JOBS = 4\n    else:\n        N_JOBS = 0\n        \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        \n    xm.master_print(f\"Train for {len(train_loader)} steps per epoch\")\n        \n    #my_model = models.densenet201(pretrained = True)\n    #my_model = pmodels.densenet201(num_classes=1000, pretrained = 'imagenet')\n    my_model = EfficientNet.from_pretrained('efficientnet-b0')\n        \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        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                torch.save({'model': model.state_dict(), \n                                'score': best_score}, \n                               f'resnet_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'resnet_fold.pth')\n\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":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":4}