{"cells":[{"metadata":{"_kg_hide-input":false,"_kg_hide-output":false},"cell_type":"markdown","source":"# Bengali.AI: Curriculum Dropout implementation\n\nThis kernel extends upon the work of iafoss in his [Grapheme fast.ai starter](https://www.kaggle.com/iafoss/grapheme-fast-ai-starter-lb-0-964) kernel which provides a strong baseline with the Bengali.AI Handwritten Grapheme Classification.\n\nHere I'm aiming to demonstrate the techniques describes in the paper [Curriculum Dropout (2017)](https://arxiv.org/abs/1703.06229) by Pietro Morerio, Jacopo Cavazza, Riccardo Volpi, Rene Vidal and Vittorio Murino.\n\nThe idea is quite simple: instead of using a fixed dropout, we instead \"start easy\" and use very little dropout in the early stages of training and gradually increase the amount of dropout until we reach our desired level.\n\n![image.png](https://snipboard.io/PUWxZj.jpg)\n\n\nI've implemented it as a [fastai v1](https://github.com/fastai/fastai) callback but hopefully the code is simple enough to be ported to any framework."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"%reload_ext autoreload\n%autoreload 2\n%matplotlib inline\n\nimport fastai\nfrom fastai.vision import *\nfrom fastai.callbacks import *\nfrom typing import Dict\nimport os\nfrom sklearn.model_selection import KFold\nfrom radam import *\nfrom csvlogger import *\nfrom mish_activation import *\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nfastai.__version__","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sz = 128\nbs = 128\nnfolds = 4 #keep the same split as the initial dataset\nfold = 0\nSEED = 2019\nTRAIN = '../input/grapheme-imgs-128x128/'\nLABELS = '../input/bengaliai-cv19/train.csv'\narch = models.densenet121\n\ndef seed_everything(seed):\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    torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Module Param Scheduler"},{"metadata":{},"cell_type":"markdown","source":"The below fastai callback was added for the purpose of dropout scheduling, though it can be used to schedule any PyTorch nn.Module parameter.\n\nIt takes in a `Learner`, as is standard with fastai callbacks, as well as the Module and parameter attribute to set. You also set the range of the attribute values (start and endpoint). Optionally, you can choose to set the epoch in which the parameter will be at max or use the epoch count passed into the learner. Lastly, the function can be passed that defines a function for scheduling the parameter, defaulting to cosign annealing by default.\n\nCosign annealing looks like this over 10 epochs with a batch size of 128 (total steps = 10 * 128 = 1280)"},{"metadata":{"trusted":true},"cell_type":"code","source":"n = 10 * 128\nplt.plot([annealing_cos(0.0000001, 0.5, i/n) for i in range(n)])\nplt.title(f'Cosigning annealing over {n} steps')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from fastai.callbacks import *\n\n\nclass ModuleParamScheduler(Callback):\n    \n    \"\"\"\n    Torch module parameter scheduler callback for fastai v1.\n    \n    Params:\n    \n      * learn - Fastai learner.\n      * module_param - Name of the parameter to fetch from the PyTorch module.\n      * module_getter - Function to fetch model layer from the PyTorch module.\n      * max_at_epoch - Epoch in which the parameter value should be at max.\n      * param_range - Start and endpoint of the parameter value.\n      * change_func - Function used to modify the value of the parameter over time.\n    \"\"\"\n\n    def __init__(\n        self,\n        learn,\n        module_param: str,\n        module_getter: Callable,\n        param_range: Tuple[float, float] = (0., 5.),\n        max_at_epoch: int=None,\n\n        change_func = annealing_cos\n    ):\n        super().__init__()\n    \n        self.learn = learn\n        self.module_param = module_param\n        self.module_getter = module_getter\n    \n        self.max_at_epoch = max_at_epoch\n        self.param_range = param_range\n        self.change_func = change_func\n    \n        self.start_epoch = None\n        self.total_epochs = None\n        \n    def on_train_begin(self, n_epochs:int, epoch:int, **kwargs:Any)->None:\n        \"Initialize our optimization params based on our annealing schedule.\"\n        res = {'epoch':self.start_epoch} if self.start_epoch is not None else None\n    \n        self.start_epoch = ifnone(self.start_epoch, epoch)\n        self.total_epochs = ifnone(self.total_epochs, self.max_at_epoch or n_epochs)\n    \n        total_steps = len(self.learn.data.train_dl) * self.total_epochs\n\n        self.scheduler = Scheduler(vals=self.param_range, n_iter=total_steps, func=self.change_func)\n\n        return res\n    \n    def on_batch_begin(self, train, **kwargs:Any)->None:\n        \"Take one step forward on the annealing schedule for the optim params.\"\n        if train:\n            self.last_val = self.scheduler.step()\n            \n            module = self.module_getter(self.learn.model)\n\n            setattr(module, self.module_param, self.last_val)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv(LABELS)\nnunique = list(df.nunique())[1:-1]\nprint(nunique)\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"stats = ([0.0692], [0.2051])\ndata = (ImageList.from_df(df, path='.', folder=TRAIN, suffix='.png', \n        cols='image_id', convert_mode='L')\n        .split_by_idx(range(fold*len(df)//nfolds,(fold+1)*len(df)//nfolds))\n        .label_from_df(cols=['grapheme_root','vowel_diacritic','consonant_diacritic'])\n        .transform(get_transforms(do_flip=False,max_warp=0.1), size=sz, padding_mode='zeros')\n        .databunch(bs=bs)).normalize(stats)\n\ndata.show_batch()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{},"cell_type":"markdown","source":"The starter model is based on DenseNet121, which I found to work quite well for such kind of problems. The first conv is replaced to accommodate for 1 channel input, and the corresponding pretrained weights are summed. ReLU activation in the head is replaced by Mish, which works noticeably better for all tasks I checked it so far. Since each portion of the prediction (grapheme_root, vowel_diacritic, and consonant_diacritic) is quite independent concept, I create a separate head for each of them (though I didn't do a comparison with one head solution)."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Head(nn.Module):\n    def __init__(self, nc, n, ps=0.5):\n        super().__init__()\n        layers = [AdaptiveConcatPool2d(), Mish(), Flatten()] + \\\n            bn_drop_lin(nc*2, 512, True, ps, Mish()) + \\\n            bn_drop_lin(512, n, True, ps)\n        self.fc = nn.Sequential(*layers)\n        self._init_weight()\n        \n    def _init_weight(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1.0)\n                m.bias.data.zero_()\n        \n    def forward(self, x):\n        return self.fc(x)\n\n#change the first conv to accept 1 chanel input\nclass Dnet_1ch(nn.Module):\n    def __init__(self, arch=arch, n=nunique, pre=True, ps=0.5):\n        super().__init__()\n        m = arch(True) if pre else arch()\n        \n        conv = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n        w = (m.features.conv0.weight.sum(1)).unsqueeze(1)\n        conv.weight = nn.Parameter(w)\n        \n        self.layer0 = nn.Sequential(conv, m.features.norm0, nn.ReLU(inplace=True))\n        self.layer1 = nn.Sequential(\n            nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False),\n            m.features.denseblock1)\n        self.layer2 = nn.Sequential(m.features.transition1,m.features.denseblock2)\n        self.layer3 = nn.Sequential(m.features.transition2,m.features.denseblock3)\n        self.layer4 = nn.Sequential(m.features.transition3,m.features.denseblock4,\n                                    m.features.norm5)\n        \n        nc = self.layer4[-1].weight.shape[0]\n        self.head1 = Head(nc,n[0])\n        self.head2 = Head(nc,n[1])\n        self.head3 = Head(nc,n[2])\n        #to_Mish(self.layer0), to_Mish(self.layer1), to_Mish(self.layer2)\n        #to_Mish(self.layer3), to_Mish(self.layer4)\n        \n    def forward(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        \n        x1 = self.head1(x)\n        x2 = self.head2(x)\n        x3 = self.head3(x)\n        \n        return x1,x2,x3","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss"},{"metadata":{},"cell_type":"markdown","source":"Cross entropy loss is applied independently to each part of the prediction and the result is summed with the corresponding weight."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Loss_combine(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n    def forward(self, input, target,reduction='mean'):\n        x1,x2,x3 = input\n        x1,x2,x3 = x1.float(),x2.float(),x3.float()\n        y = target.long()\n        return 0.7*F.cross_entropy(x1,y[:,0],reduction=reduction) + 0.1*F.cross_entropy(x2,y[:,1],reduction=reduction) + \\\n          0.2*F.cross_entropy(x3,y[:,2],reduction=reduction)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The code below computes the competition metric and recall macro metrics for individual components of the prediction. The code is partially borrowed from fast.ai."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class Metric_idx(Callback):\n    def __init__(self, idx, average='macro'):\n        super().__init__()\n        self.idx = idx\n        self.n_classes = 0\n        self.average = average\n        self.cm = None\n        self.eps = 1e-9\n        \n    def on_epoch_begin(self, **kwargs):\n        self.tp = 0\n        self.fp = 0\n        self.cm = None\n    \n    def on_batch_end(self, last_output:Tensor, last_target:Tensor, **kwargs):\n        last_output = last_output[self.idx]\n        last_target = last_target[:,self.idx]\n        preds = last_output.argmax(-1).view(-1).cpu()\n        targs = last_target.long().cpu()\n        \n        if self.n_classes == 0:\n            self.n_classes = last_output.shape[-1]\n            self.x = torch.arange(0, self.n_classes)\n        cm = ((preds==self.x[:, None]) & (targs==self.x[:, None, None])) \\\n          .sum(dim=2, dtype=torch.float32)\n        if self.cm is None: self.cm =  cm\n        else:               self.cm += cm\n\n    def _weights(self, avg:str):\n        if self.n_classes != 2 and avg == \"binary\":\n            avg = self.average = \"macro\"\n            warn(\"average=`binary` was selected for a non binary case. \\\n                 Value for average has now been set to `macro` instead.\")\n        if avg == \"binary\":\n            if self.pos_label not in (0, 1):\n                self.pos_label = 1\n                warn(\"Invalid value for pos_label. It has now been set to 1.\")\n            if self.pos_label == 1: return Tensor([0,1])\n            else: return Tensor([1,0])\n        elif avg == \"micro\": return self.cm.sum(dim=0) / self.cm.sum()\n        elif avg == \"macro\": return torch.ones((self.n_classes,)) / self.n_classes\n        elif avg == \"weighted\": return self.cm.sum(dim=1) / self.cm.sum()\n        \n    def _recall(self):\n        rec = torch.diag(self.cm) / (self.cm.sum(dim=1) + self.eps)\n        if self.average is None: return rec\n        else:\n            if self.average == \"micro\": weights = self._weights(avg=\"weighted\")\n            else: weights = self._weights(avg=self.average)\n            return (rec * weights).sum()\n    \n    def on_epoch_end(self, last_metrics, **kwargs): \n        return add_metrics(last_metrics, self._recall())\n    \nMetric_grapheme = partial(Metric_idx,0)\nMetric_vowel = partial(Metric_idx,1)\nMetric_consonant = partial(Metric_idx,2)\n\nclass Metric_tot(Callback):\n    def __init__(self):\n        super().__init__()\n        self.grapheme = Metric_idx(0)\n        self.vowel = Metric_idx(1)\n        self.consonant = Metric_idx(2)\n        \n    def on_epoch_begin(self, **kwargs):\n        self.grapheme.on_epoch_begin(**kwargs)\n        self.vowel.on_epoch_begin(**kwargs)\n        self.consonant.on_epoch_begin(**kwargs)\n    \n    def on_batch_end(self, last_output:Tensor, last_target:Tensor, **kwargs):\n        self.grapheme.on_batch_end(last_output, last_target, **kwargs)\n        self.vowel.on_batch_end(last_output, last_target, **kwargs)\n        self.consonant.on_batch_end(last_output, last_target, **kwargs)\n        \n    def on_epoch_end(self, last_metrics, **kwargs): \n        return add_metrics(last_metrics, 0.5*self.grapheme._recall() +\n                0.25*self.vowel._recall() + 0.25*self.consonant._recall())","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":false,"trusted":true},"cell_type":"code","source":"#fix the issue in fast.ai of saving gradients along with weights\n#so only weights are written, and files are ~4 times smaller\n\nclass SaveModelCallback(TrackerCallback):\n    \"A `TrackerCallback` that saves the model when monitored quantity is best.\"\n    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto',\n                 every:str='improvement', name:str='bestmodel'):\n        super().__init__(learn, monitor=monitor, mode=mode)\n        self.every,self.name = every,name\n        if self.every not in ['improvement', 'epoch']:\n            warn(f'SaveModel every {self.every} is invalid, falling back to \"improvement\".')\n            self.every = 'improvement'\n                 \n    def jump_to_epoch(self, epoch:int)->None:\n        try: \n            self.learn.load(f'{self.name}_{epoch-1}', purge=False)\n            print(f\"Loaded {self.name}_{epoch-1}\")\n        except: print(f'Model {self.name}_{epoch-1} not found.')\n\n    def on_epoch_end(self, epoch:int, **kwargs:Any)->None:\n        \"Compare the value monitored to its best score and maybe save the model.\"\n        if self.every==\"epoch\": \n            #self.learn.save(f'{self.name}_{epoch}')\n            torch.save(learn.model.state_dict(),f'{self.name}_{epoch}.pth')\n        else: #every=\"improvement\"\n            current = self.get_monitor_value()\n            if current is not None and self.operator(current, self.best):\n                #print(f'Better model found at epoch {epoch} \\\n                #  with {self.monitor} value: {current}.')\n                self.best = current\n                #self.learn.save(f'{self.name}')\n                torch.save(learn.model.state_dict(),f'{self.name}.pth')\n\n    def on_train_end(self, **kwargs):\n        \"Load the best model.\"\n        if self.every==\"improvement\" and os.path.isfile(f'{self.name}.pth'):\n            #self.learn.load(f'{self.name}', purge=False)\n            self.model.load_state_dict(torch.load(f'{self.name}.pth'))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## MixUp"},{"metadata":{},"cell_type":"markdown","source":"The code below modifies fast.ai MixUp calback to make it compatible with the current data."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class MixUpLoss(Module):\n    \"Adapt the loss function `crit` to go with mixup.\"\n    \n    def __init__(self, crit, reduction='mean'):\n        super().__init__()\n        if hasattr(crit, 'reduction'): \n            self.crit = crit\n            self.old_red = crit.reduction\n            setattr(self.crit, 'reduction', 'none')\n        else: \n            self.crit = partial(crit, reduction='none')\n            self.old_crit = crit\n        self.reduction = reduction\n        \n    def forward(self, output, target):\n        if len(target.shape) == 2 and target.shape[1] == 7:\n            loss1, loss2 = self.crit(output,target[:,0:3].long()), self.crit(output,target[:,3:6].long())\n            d = loss1 * target[:,-1] + loss2 * (1-target[:,-1])\n        else:  d = self.crit(output, target)\n        if self.reduction == 'mean':    return d.mean()\n        elif self.reduction == 'sum':   return d.sum()\n        return d\n    \n    def get_old(self):\n        if hasattr(self, 'old_crit'):  return self.old_crit\n        elif hasattr(self, 'old_red'): \n            setattr(self.crit, 'reduction', self.old_red)\n            return self.crit\n\nclass MixUpCallback(LearnerCallback):\n    \"Callback that creates the mixed-up input and target.\"\n    def __init__(self, learn:Learner, alpha:float=0.4, stack_x:bool=False, stack_y:bool=True):\n        super().__init__(learn)\n        self.alpha,self.stack_x,self.stack_y = alpha,stack_x,stack_y\n    \n    def on_train_begin(self, **kwargs):\n        if self.stack_y: self.learn.loss_func = MixUpLoss(self.learn.loss_func)\n        \n    def on_batch_begin(self, last_input, last_target, train, **kwargs):\n        \"Applies mixup to `last_input` and `last_target` if `train`.\"\n        if not train: return\n        lambd = np.random.beta(self.alpha, self.alpha, last_target.size(0))\n        lambd = np.concatenate([lambd[:,None], 1-lambd[:,None]], 1).max(1)\n        lambd = last_input.new(lambd)\n        shuffle = torch.randperm(last_target.size(0)).to(last_input.device)\n        x1, y1 = last_input[shuffle], last_target[shuffle]\n        if self.stack_x:\n            new_input = [last_input, last_input[shuffle], lambd]\n        else: \n            out_shape = [lambd.size(0)] + [1 for _ in range(len(x1.shape) - 1)]\n            new_input = (last_input * lambd.view(out_shape) + x1 * (1-lambd).view(out_shape))\n        if self.stack_y:\n            new_target = torch.cat([last_target.float(), y1.float(), lambd[:,None].float()], 1)\n        else:\n            if len(last_target.shape) == 2:\n                lambd = lambd.unsqueeze(1).float()\n            new_target = last_target.float() * lambd + y1.float() * (1-lambd)\n        return {'last_input': new_input, 'last_target': new_target}  \n    \n    def on_train_end(self, **kwargs):\n        if self.stack_y: self.learn.loss_func = self.learn.loss_func.get_old()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{},"cell_type":"markdown","source":"I have performed a check of different optimizers and schedules on a [similar task](https://www.kaggle.com/c/Kannada-MNIST/discussion/122430), and [Over9000 optimizer](https://github.com/mgrankin/over9000) cosine annealing **without warm-up** worked the best. Freezing the backbone at the initial stage of training didn't give me any advantage in that test, so here I perform the training straight a way with discriminative learning rate (smaller lr for backbone)."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"model = Dnet_1ch()\nlearn = Learner(data, model, loss_func=Loss_combine(), opt_func=Over9000,\n        metrics=[Metric_grapheme(),Metric_vowel(),Metric_consonant(),Metric_tot()])\nlogger = CSVLogger(learn,f'log{fold}')\nlearn.clip_grad = 1.0\nlearn.split([model.head1])\nlearn.unfreeze()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dropout Scheduling\n\nI create a bunch of different schedulers each of them operating on a different dropout layer. This can obviously be customised however you'd like and can ever be applied to other model hyperparameters."},{"metadata":{"trusted":true},"cell_type":"code","source":"dropout_sched_callbacks = [\n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head1.fc[4], param_range=(0., 0.5)),\n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head1.fc[8], param_range=(0., 0.5)),\n    \n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head2.fc[4], param_range=(0., 0.5)),\n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head2.fc[8], param_range=(0., 0.5)),\n    \n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head3.fc[4], param_range=(0., 0.5)),\n    ModuleParamScheduler(learn=learn, module_param='p', module_getter=lambda model: model.head3.fc[8], param_range=(0., 0.5))\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(\n    32,\n    max_lr=slice(0.2e-2, 1e-2),\n    wd=[1e-3, 0.1e-1],\n    pct_start=0.0, \n    div_factor=100,\n    callbacks=[\n        logger,\n        SaveModelCallback(learn, monitor='metric_tot', mode='max',name=f'model_{fold}'),\n        MixUpCallback(learn)\n    ] + dropout_sched_callbacks\n)\n\n# metrics: Metric_grapheme, Metric_vowel, Metric_consonant, Metric_tot (competition metric)","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":1}