{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"tpu = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if tpu:\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 \"nightly\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install efficientnet_pytorch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.optim as optim\nfrom torchvision import transforms\nfrom torch.optim import lr_scheduler\nfrom efficientnet_pytorch import EfficientNet\nfrom torch.nn import functional as F\n\nif tpu:\n    import torch_xla\n    import torch_xla.debug.metrics as met\n    import torch_xla.distributed.data_parallel as dp\n    import torch_xla.distributed.parallel_loader as pl\n\n    import torch_xla.utils.utils as xu\n    import torch_xla.core.xla_model as xm\n    import torch_xla.utils.serialization as xser\n    import torch_xla.distributed.xla_multiprocessing as xmp\n    import torch_xla.test.test_utils as test_utils\n\nfrom warmup_scheduler import GradualWarmupScheduler\n\nimport numpy as np\nimport pandas as pd\nimport os\nimport gc\nimport time\nfrom tqdm import tqdm\nimport timm\n\nimport cv2\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import StratifiedKFold\nimport albumentations\nfrom albumentations import Compose, Normalize, HorizontalFlip, VerticalFlip\n\nfrom albumentations import (\n    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize\n)\nfrom albumentations.pytorch import ToTensorV2\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(torch.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#os.environ[\"XLA_USE_BF16\"] = \"1\"\nif tpu:\n    os.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#config\n\n\ndata_dir = '../input/cassava-leaf-disease-classification'\ndf_train = pd.read_csv(os.path.join(data_dir, 'train.csv'))\nimage_folder = os.path.join(data_dir, 'train_images')\n\n\nfolds = 5\npresent_fold = 2\nimage_size = 384\nbatch_size = 8\nval_batch_size = 8\nnum_workers = 4\nout_dim = 5\ninit_lr = 1e-4\n\nwarmup_factor = 7\nwarmup_epo = 1\n\nsmoothing = 0.05\n\nepoch_threshold = 5\n\ndebug = False\n\nn_epochs = 2 if debug else 3\n\nkernel_type = 'vit_large_patch16_384' \n\nnet_type = 'vit_large_patch16_384'\nprint(image_folder)\nt1 = 0.8\nt2 = 1.4","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#splitting data\nskf = StratifiedKFold(folds,shuffle=True,random_state=42)\ndf_train['fold'] = -1\nif debug:\n    df_train = df_train[:500]\nfor i, (train_idx, valid_idx) in enumerate(skf.split(df_train, df_train['label'])):\n    df_train.loc[valid_idx, 'fold'] = i\ndf_train.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CASSAVA(Dataset):\n    def __init__(self,df,transforms=None):\n        self.df = df\n        #print(self.df.shape[0])\n        self.transforms = transforms\n    def __len__(self):\n        return self.df.shape[0]\n    def __getitem__(self,idx):\n        row = self.df.iloc[idx]\n        img_dir = os.path.join(image_folder,row['image_id'])\n        img = cv2.imread(img_dir)\n        if self.transforms is not None:\n            img = self.transforms(image=img)['image']\n        label = row.label\n        return img,torch.tensor(label)\n#dummy\ndataset = CASSAVA(df_train)\nX,Y = dataset.__getitem__(3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return Compose([\n            RandomResizedCrop(image_size, image_size),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(p=0.5),\n            HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n            RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            CoarseDropout(p=0.5),\n            Cutout(p=0.5),\n            ToTensorV2(p=1.0),\n        ], p=1.)\n  \n        \ndef get_valid_transforms():\n    return Compose([\n            #HorizontalFlip(p=0.5),\n            #VerticalFlip(p=0.5),\n            CenterCrop(image_size, image_size, p=1.),\n            Resize(image_size, image_size),\n            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.)\ntransforms_train = get_train_transforms()\ntransforms_valid = get_valid_transforms()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_show = CASSAVA(df_train, transforms=transforms_train)\n\nfor i in range(2):\n    f, axarr = plt.subplots(1,5,figsize=(20,10))\n    for p in range(5):\n        idx = np.random.randint(0, len(dataset_show))\n        img, label = dataset_show[idx]\n        print(type(img),img.size(),img.max())\n        axarr[p].imshow(img.transpose(0, 1).transpose(1,2).squeeze())\n        axarr[p].set_title(str(torch.argmax(label,axis=-1)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Classifier(nn.Module):\n    def __init__(self,):\n        super().__init__()\n        self.effn = EfficientNet.from_pretrained('efficientnet-b4')\n       \n        self.dense = nn.Linear(1792,out_dim)\n        self.pooling = nn.MaxPool2d((12,12))\n        self.flatten = nn.Flatten()\n        for param in self.effn.parameters():\n            param.requires_grad = True\n        \n    def forward(self,x):\n        x = self.effn.extract_features(x)\n        x = self.pooling(x)\n        x = self.flatten(x)\n        x = self.dense(x)\n       \n        return x\n    \nclass cassavamodel(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        #print(self.model)\n        n_features = self.model.head.in_features\n        self.model.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n    \n#dummy\n'''dummy_img = torch.rand((1,3,image_size,image_size))\ndummy = Classifier()\n\n#print(dummy)\no_p = dummy(dummy_img)\nprint(o_p)\nmodel = Classifier()\ndummy_img = torch.rand((1,3,image_size,image_size))\no_p = model(dummy_img)\nprint(o_p)'''","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TaylorSoftmax(nn.Module):\n\n    def __init__(self, dim=1, n=2):\n        super(TaylorSoftmax, self).__init__()\n        assert n % 2 == 0\n        self.dim = dim\n        self.n = n\n\n    def forward(self, x):\n        \n        fn = torch.ones_like(x)\n       \n        denor = 1.\n        for i in range(1, self.n+1):\n            denor *= i\n            fn = fn + x.pow(i) /(denor+1e-6)\n\n        out = fn / fn.sum(dim=self.dim, keepdims=True)\n        return out\n    \nclass LabelSmoothingLoss(nn.Module):\n\n    def __init__(self, classes, smoothing=0.0, dim=-1): \n        super(LabelSmoothingLoss, self).__init__() \n        self.confidence = 1.0 - smoothing \n        self.smoothing = smoothing \n        self.cls = classes \n        self.dim = dim \n    def forward(self, pred, target): \n        \"\"\"Taylor Softmax and log are already applied on the logits\"\"\"\n        with torch.no_grad(): \n            true_dist = torch.zeros_like(pred) \n            if self.cls-1==0:\n                raise Exception('self.cls = 1')\n            true_dist.fill_(self.smoothing / (self.cls - 1)) \n            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) \n        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))\n    \n\nclass TaylorCrossEntropyLoss(nn.Module):\n\n    def __init__(self, n=2, ignore_index=-1, reduction='mean', smoothing=0.2):\n        super(TaylorCrossEntropyLoss, self).__init__()\n        assert n % 2 == 0\n        self.taylor_softmax = TaylorSoftmax(dim=-1, n=n)\n        self.reduction = reduction\n        self.ignore_index = ignore_index\n        self.lab_smooth = LabelSmoothingLoss(out_dim, smoothing=smoothing)\n\n    def forward(self, logits, labels):\n \n        log_probs = self.taylor_softmax(logits).log()\n\n        loss = self.lab_smooth(log_probs, labels)\n        return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def log_t(u, t):\n    \"\"\"Compute log_t for `u'.\"\"\"\n    if t==1.0:\n        return u.log()\n    else:\n        return (u.pow(1.0 - t) - 1.0) / (1.0 - t)\n\ndef exp_t(u, t):\n    \"\"\"Compute exp_t for `u'.\"\"\"\n    if t==1:\n        return u.exp()\n    else:\n        return (1.0 + (1.0-t)*u).relu().pow(1.0 / (1.0 - t))\n\ndef compute_normalization_fixed_point(activations, t, num_iters):\n\n    \"\"\"Returns the normalization value for each example (t > 1.0).\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (> 1.0 for tail heaviness).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same shape as activation with the last dimension being 1.\n    \"\"\"\n    mu, _ = torch.max(activations, -1, keepdim=True)\n    normalized_activations_step_0 = activations - mu\n\n    normalized_activations = normalized_activations_step_0\n\n    for _ in range(num_iters):\n        logt_partition = torch.sum(\n                exp_t(normalized_activations, t), -1, keepdim=True)\n        normalized_activations = normalized_activations_step_0 * \\\n                logt_partition.pow(1.0-t)\n\n    logt_partition = torch.sum(\n            exp_t(normalized_activations, t), -1, keepdim=True)\n    normalization_constants = - log_t(1.0 / logt_partition, t) + mu\n\n    return normalization_constants\n\ndef compute_normalization_binary_search(activations, t, num_iters):\n\n    \"\"\"Returns the normalization value for each example (t < 1.0).\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (< 1.0 for finite support).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same rank as activation with the last dimension being 1.\n    \"\"\"\n\n    mu, _ = torch.max(activations, -1, keepdim=True)\n    normalized_activations = activations - mu\n\n    effective_dim = \\\n        torch.sum(\n                (normalized_activations > -1.0 / (1.0-t)).to(torch.int32),\n            dim=-1, keepdim=True).to(activations.dtype)\n\n    shape_partition = activations.shape[:-1] + (1,)\n    lower = torch.zeros(shape_partition, dtype=activations.dtype, device=activations.device)\n    upper = -log_t(1.0/effective_dim, t) * torch.ones_like(lower)\n\n    for _ in range(num_iters):\n        logt_partition = (upper + lower)/2.0\n        sum_probs = torch.sum(\n                exp_t(normalized_activations - logt_partition, t),\n                dim=-1, keepdim=True)\n        update = (sum_probs < 1.0).to(activations.dtype)\n        lower = torch.reshape(\n                lower * update + (1.0-update) * logt_partition,\n                shape_partition)\n        upper = torch.reshape(\n                upper * (1.0 - update) + update * logt_partition,\n                shape_partition)\n\n    logt_partition = (upper + lower)/2.0\n    return logt_partition + mu\n\nclass ComputeNormalization(torch.autograd.Function):\n    \"\"\"\n    Class implementing custom backward pass for compute_normalization. See compute_normalization.\n    \"\"\"\n    @staticmethod\n    def forward(ctx, activations, t, num_iters):\n        if t < 1.0:\n            normalization_constants = compute_normalization_binary_search(activations, t, num_iters)\n        else:\n            normalization_constants = compute_normalization_fixed_point(activations, t, num_iters)\n\n        ctx.save_for_backward(activations, normalization_constants)\n        ctx.t=t\n        return normalization_constants\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        activations, normalization_constants = ctx.saved_tensors\n        t = ctx.t\n        normalized_activations = activations - normalization_constants \n        probabilities = exp_t(normalized_activations, t)\n        escorts = probabilities.pow(t)\n        escorts = escorts / escorts.sum(dim=-1, keepdim=True)\n        grad_input = escorts * grad_output\n        \n        return grad_input, None, None\n\ndef compute_normalization(activations, t, num_iters=5):\n    \"\"\"Returns the normalization value for each example. \n    Backward pass is implemented.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same rank as activation with the last dimension being 1.\n    \"\"\"\n    return ComputeNormalization.apply(activations, t, num_iters)\n\ndef tempered_sigmoid(activations, t, num_iters = 5):\n    \"\"\"Tempered sigmoid function.\n    Args:\n      activations: Activations for the positive class for binary classification.\n      t: Temperature tensor > 0.0.\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A probabilities tensor.\n    \"\"\"\n    internal_activations = torch.stack([activations,\n        torch.zeros_like(activations)],\n        dim=-1)\n    internal_probabilities = tempered_softmax(internal_activations, t, num_iters)\n    return internal_probabilities[..., 0]\n\n\ndef tempered_softmax(activations, t, num_iters=5):\n    \"\"\"Tempered softmax function.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature > 1.0.\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A probabilities tensor.\n    \"\"\"\n    if t == 1.0:\n        return activations.softmax(dim=-1)\n\n    normalization_constants = compute_normalization(activations, t, num_iters)\n    return exp_t(activations - normalization_constants, t)\n\ndef bi_tempered_binary_logistic_loss(activations,\n        labels,\n        t1,\n        t2,\n        label_smoothing = 0.0,\n        num_iters=5,\n        reduction='mean'):\n\n    \"\"\"Bi-Tempered binary logistic loss.\n    Args:\n      activations: A tensor containing activations for class 1.\n      labels: A tensor with shape as activations, containing probabilities for class 1\n      t1: Temperature 1 (< 1.0 for boundedness).\n      t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      label_smoothing: Label smoothing\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A loss tensor.\n    \"\"\"\n    internal_activations = torch.stack([activations,\n        torch.zeros_like(activations)],\n        dim=-1)\n    internal_labels = torch.stack([labels.to(activations.dtype),\n        1.0 - labels.to(activations.dtype)],\n        dim=-1)\n    return bi_tempered_logistic_loss(internal_activations, \n            internal_labels,\n            t1,\n            t2,\n            label_smoothing = label_smoothing,\n            num_iters = num_iters,\n            reduction = reduction)\n\ndef bi_tempered_logistic_loss(activations,\n        labels,\n        t1,\n        t2,\n        label_smoothing=0.0,\n        num_iters=5,\n        reduction = 'mean'):\n\n    \"\"\"Bi-Tempered Logistic Loss.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      labels: A tensor with shape and dtype as activations (onehot), \n        or a long tensor of one dimension less than activations (pytorch standard)\n      t1: Temperature 1 (< 1.0 for boundedness).\n      t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      label_smoothing: Label smoothing parameter between [0, 1). Default 0.0.\n      num_iters: Number of iterations to run the method. Default 5.\n      reduction: ``'none'`` | ``'mean'`` | ``'sum'``. Default ``'mean'``.\n        ``'none'``: No reduction is applied, return shape is shape of\n        activations without the last dimension.\n        ``'mean'``: Loss is averaged over minibatch. Return shape (1,)\n        ``'sum'``: Loss is summed over minibatch. Return shape (1,)\n    Returns:\n      A loss tensor.\n    \"\"\"\n\n    if len(labels.shape)<len(activations.shape): #not one-hot\n        labels_onehot = torch.zeros_like(activations)\n        labels_onehot.scatter_(1, labels[..., None], 1)\n    else:\n        labels_onehot = labels\n\n    if label_smoothing > 0:\n        num_classes = labels_onehot.shape[-1]\n        labels_onehot = ( 1 - label_smoothing * num_classes / (num_classes - 1) ) \\\n                * labels_onehot + \\\n                label_smoothing / (num_classes - 1)\n\n    probabilities = tempered_softmax(activations, t2, num_iters)\n\n    loss_values = labels_onehot * log_t(labels_onehot + 1e-10, t1) \\\n            - labels_onehot * log_t(probabilities, t1) \\\n            - labels_onehot.pow(2.0 - t1) / (2.0 - t1) \\\n            + probabilities.pow(2.0 - t1) / (2.0 - t1)\n    loss_values = loss_values.sum(dim = -1) #sum over classes\n\n    if reduction == 'none':\n        return loss_values\n    if reduction == 'sum':\n        return loss_values.sum()\n    if reduction == 'mean':\n        return loss_values.mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def reduce_fn(vals):\n    # take average\n    return sum(vals) / len(vals)\ndef train_one_epoch(model,device,loader,optimizer):\n    model.train()\n    train_loss = []\n    bar = tqdm(loader)\n    for image,label in bar:\n        image = image.type(torch.FloatTensor)\n        image = image.to(device)\n        label = label.to(device)\n    \n        optimizer.zero_grad()\n        y_pred = model(image)\n        loss = bi_tempered_logistic_loss(y_pred, label, t1=t1, t2=t2, label_smoothing=smoothing)\n\n        loss.backward()\n        if tpu:\n            xm.optimizer_step(optimizer)\n        else:\n            optimizer.step()\n        loss_np = loss.detach().cpu().numpy()\n        #xm.master_print(loss_np)\n        train_loss.append(loss_np)\n    return train_loss\n\ndef validate_one_epoch(model,device,loader):\n    model.eval()\n    val_loss = []\n    sample_num = 0\n    y_preds = []\n    y_trues = []\n    \n    bar = tqdm(loader)\n    for image,label in bar:\n        image = image.type(torch.FloatTensor)\n        image = image.to(device)\n        label = label.to(device)\n        pred = model(image)\n        loss = bi_tempered_logistic_loss(pred,label, t1=t1, t2=t2, label_smoothing=smoothing)\n        \n        pred_ = torch.argmax(pred,dim=-1)\n        y_preds+=list(pred_.detach().cpu().numpy())\n        \n        y_trues+=list(label.cpu().numpy())\n        \n        \n        val_loss.append(loss.detach().cpu().numpy())\n        \n        sample_num += image.shape[0]\n    val_loss = np.mean(val_loss)\n    y_preds = np.asarray(y_preds)\n    y_trues = np.asarray(y_trues)\n    acc = np.asarray(y_trues==y_preds,dtype=np.float32)\n    if tpu:\n        accuracy = xm.mesh_reduce('test_accuracy', acc, np.mean)\n        xm.master_print(\"Validation Accuracy = \",accuracy)\n    else:\n        accuracy = np.mean(acc)\n        print('Validation Accuracy = ',accuracy)\n    return val_loss,accuracy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#model = cassavamodel(net_type, n_class=out_dim)\nmodel = Classifier()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GradualWarmupSchedulerV2(GradualWarmupScheduler):\n    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):\n        super(GradualWarmupSchedulerV2, self).__init__(optimizer, multiplier, total_epoch, after_scheduler)\n    def get_lr(self):\n        if self.last_epoch > self.total_epoch:\n            if self.after_scheduler:\n                if not self.finished:\n                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]\n                    self.finished = True\n                return self.after_scheduler.get_lr()\n            return [base_lr * self.multiplier for base_lr in self.base_lrs]\n        if self.multiplier == 1.0:\n            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]\n        else:\n            return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]\n\nif tpu:\n    init_lr = init_lr * xm.xrt_world_size()\n\noptimizer = optim.Adam(model.parameters(), lr=init_lr)\nscheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs-warmup_epo)\nscheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n\n\n\nlrs = []\nfor epoch in range(1, n_epochs+1):\n    scheduler_warmup.step(epoch-1)\n    lrs.append(optimizer.param_groups[0][\"lr\"])\nplt.figure(figsize=(20,3))\nplt.plot(lrs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_model(model=model):\n    for fold in range(folds):\n        if fold!=present_fold:\n            continue\n        if tpu:\n            device = xm.xla_device()\n        else:\n            device = 'cuda'\n        model = model.to(device)\n        model.load_state_dict(torch.load('../input/effnetb4/vit_large_patch16_384_best_fold2.pth'))\n        optimizer =  optim.Adam(model.parameters(), lr=init_lr/warmup_factor)\n\n        train_idx = np.where((df_train['fold'] != fold))[0]\n        valid_idx = np.where((df_train['fold'] == fold))[0]\n        \n        df_curr = df_train.loc[train_idx]\n        df_val = df_train.loc[valid_idx]\n        #print(df_curr.shape)\n        \n        scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs-warmup_epo)\n        scheduler = GradualWarmupScheduler(optimizer, multiplier=warmup_factor, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n\n\n        dataset_train = CASSAVA(df_curr , transforms=transforms_train)\n        dataset_valid = CASSAVA(df_val, transforms=transforms_valid)\n        \n        if tpu:\n\n            train_sampler = torch.utils.data.distributed.DistributedSampler(\n              dataset_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            dataset=dataset_train,\n            batch_size=batch_size,\n            sampler=train_sampler if tpu else None,\n            drop_last=True,\n            num_workers=num_workers\n        )\n        if tpu:    \n            valid_sampler = torch.utils.data.distributed.DistributedSampler(\n              dataset_valid,\n              num_replicas=xm.xrt_world_size(),\n              rank=xm.get_ordinal(),\n              shuffle=False)\n\n\n        valid_loader = torch.utils.data.DataLoader(\n            dataset=dataset_valid, \n            batch_size=val_batch_size,\n            sampler=valid_sampler if tpu else None,\n            drop_last=True,\n            num_workers=num_workers,\n        )\n\n        acc_max = 0.\n        \n        for epoch in range(1, n_epochs+1):\n            \n       \n#            if epoch < 4:\n#                # freeze backbone layers\n#                for param in model.children():\n#                    count +=1\n#                    if count < 4: #freezing first 3 layers\n#                        param.requires_grad = False\n\n#            else:\n#                scheduler.step(epoch-1)\n#                for param in model.children():\n#                    param.requires_grad = True'''\n             \n    \n            \n            if tpu:\n                train_loader = pl.ParallelLoader(train_loader, [device]).per_device_loader(device)\n            train_loss = train_one_epoch(model,device,train_loader, optimizer)\n    \n            gc.collect()\n            \n            if tpu:\n                valid_loader = pl.ParallelLoader(valid_loader, [device]).per_device_loader(device)\n\n            val_loss, acc = validate_one_epoch(model,device,valid_loader)\n            \n            gc.collect()\n            \n            content = time.ctime() + ' ' + f'FOLD -> {fold} --> Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, train loss: {np.mean(train_loss):.5f}, val loss: {np.mean(val_loss):.5f}, acc: {(acc):.5f}'\n\n            with open(f'log_{kernel_type}.txt', 'a') as appender:\n                appender.write(content + '\\n')\n            \n\n        \n            best_file = f'{kernel_type}_best_fold{fold}.pth'\n\n            if acc > acc_max:\n                if tpu:\n                    xm.save(model.state_dict(), os.path.join(best_file))\n                else:\n                    torch.save(model.state_dict(),os.path.join(best_file))\n                #xser.save(model.state_dict(), os.path.join(best_file))\n                acc_max = acc\n            ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ndef _mp_fn(rank, flags):\n    global acc_list\n    torch.set_default_tensor_type('torch.FloatTensor')\n    res = train_model()\nif tpu:\n    FLAGS={}\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')\nelse:\n    _mp_fn(None,None)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"f = open(f'./log_{kernel_type}.txt', \"r\")\n\nprint(f.read())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}