{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q timm\n!pip install -q torchcontrib\n!pip install -q pytorch_lightning==1.2.5\n!pip uninstall -q -y albumentations && pip install -q git+https://github.com/albumentations-team/albumentations\n# !pip install -q wandb","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\nimport numpy as np\nimport pytorch_lightning as pl\nimport torch\nimport pandas as pd\nimport timm\nimport torch.nn as nn\n\nfrom PIL import Image\nfrom sklearn.model_selection import KFold\nfrom torch.utils.data import Dataset, DataLoader\nfrom pytorch_lightning import Trainer, seed_everything\nfrom pytorch_lightning.loggers import CSVLogger\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, StochasticWeightAveraging\nfrom pytorch_lightning.metrics import Metric\nfrom torchcontrib.optim import SWA\nfrom typing import List, Dict\n\n#Augmentation\nimport albumentations as A\nfrom torchvision import transforms as tsfm\nfrom albumentations.pytorch import ToTensorV2\n\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(pl.__version__)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"code","source":"class CFG:\n    seed = 42\n    # Một số đường dẫn\n    root_dir_origin = \"../input/plant-pathology-2021-fgvc8/\"\n    root_dir_resized = \"../input/resized-plantpathology2021fgvc8-train-data-new/resized_plant-pathology-2021-fgvc8_train_data\"\n    train_csv_path = os.path.join(root_dir_origin, 'train.csv')\n    train_imgs_dir = os.path.join(root_dir_resized, 'resized_train_images_360_512')\n\n    #     folds_csv_path = \"../input/pp2021-dataset-gnueih/6folds_pp2021.csv\"\n    folds_csv_path = \"../input/pp2021-kfold-tfrecords-0/folds.csv\"\n    \n    num_classes = 5\n    labels = np.array(['powdery_mildew',\n                     'scab',\n                     'complex',\n                     'frog_eye_leaf_spot',\n                     'rust',])    \n    \n    # Version cho logger và tên model sử dụng\n    version = 'kag_final_v18'\n    model_name = 'tf_efficientnet_b4_ns'\n    \n    # Các tham số training\n    use_sgd=False #Sử dụng Adam cho tốc độ hội tụ nhanh hơn và do chỉ train 3 epoch \n    fl_alpha = 1.0  # tham số scale focal_loss - alpha\n    fl_gamma = 2.0  # tham số điều chỉnh dạng (đường có dạng hàm exponential) - focal loss\n    # Trọng số loss của từng class, tham khảo từ https://www.kaggle.com/crissallan/pytorchlightning-efficientnet-focalloss-inference\n    cls_weight = [0.3648, 0.0813, 0.2184, 0.1066, 0.2290] \n    # Trọng số positive để cân bằng tỉ lệ nhãn (số lượng sample có bệnh a nhỏ hơn nhiều so với không có bệnh)\n    pos_weight = [1.6990, 1.3010, 1.6990, 1.4771, 1.6990] \n    \n    num_epochs = 3\n    batch_size = 8\n#     scheduler_freq = 1548//36\n#     t_max = 36\n    scheduler_freq = 1855//35 # Chia một epoch làm 35 lần cập nhật learning rate\n    t_max = 35 # Set một cycle learning rate scheduler ứng với một epoch\n    lr = 3e-4 # Learning rate khởi tạo cho optimizer\n    min_lr = 1e-7 # Learning rate tối thiểu dùng CosineAnnealingWarmRestarts\n    \n    n_train_fold = 5 # Số fold dùng cho training\n    reserve_fold = 5 # Fold index không dùng cho training\n    \n    num_workers = 4 \n    accum_grad_batch = 1\n    early_stop_delta = 1e-7\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Seed cho phép reproduction","metadata":{}},{"cell_type":"code","source":"seed_everything(CFG.seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data","metadata":{}},{"cell_type":"code","source":"# Dataset Class\nclass ImageDataset(Dataset):\n    \"\"\" Leaf Disease Dataset \"\"\"\n    def __init__(self,\n                image_names,\n                labels,\n                image_dir, \n                transforms):        \n        self.image_names = image_names\n        self.image_dir = image_dir\n        self.transforms = transforms                \n        self.labels = labels\n\n    def __len__(self) -> int:\n        return len(self.image_names)\n\n    def __getitem__(self, idx: int):\n        image_path = os.path.join(self.image_dir, self.image_names[idx])   \n        image = Image.open(image_path).convert('RGB')\n        # image = np.array(Image.open(image_path))\n        target = self.labels[idx]\n        if self.transforms is not None:\n          # image = np.array(image)\n#             image = self.transforms(image=np.array(image))['image']\n            image = self.transforms(image)\n        return image, target","metadata":{"_kg_hide-input":false,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Pytorch Lightning Data Module \nclass ImageDataModule(pl.LightningDataModule):\n    def __init__(self,\n                 df: pd.DataFrame,\n                 train_transforms,\n                 valid_transforms,\n                 no_aug_transforms,\n                 image_dir: str,\n                 fold_num: int,\n                 configurations: Dict[str, int]):\n        super().__init__()\n        self.df = df\n        self.train_transforms = train_transforms\n        self.valid_transforms = valid_transforms\n        self.no_aug_transforms = no_aug_transforms\n        self.image_dir = image_dir\n        self.fold_num = fold_num\n        self.CFG = configurations\n    \n    def setup(self, stage=None) -> None:\n        train_df = self.df[df.fold != self.fold_num].reset_index()\n        valid_df = self.df[df.fold == self.fold_num].reset_index()\n        \n        print(f\"Size of Train Dataset: {len(train_df.index)}\")\n        print(f\"Size of Validation Dataset: {len(valid_df.index)}\")\n        if stage is None or stage == 'fit':\n            self.train_dataset = ImageDataset(image_names=train_df.image.values, \n                                            labels=train_df[self.CFG.labels].values, \n                                            image_dir=self.image_dir, \n                                            transforms=self.train_transforms,\n                                            )\n\n            self.valid_dataset = ImageDataset(image_names=valid_df.image.values, \n                                            labels=valid_df[self.CFG.labels].values, \n                                            image_dir=self.image_dir, \n                                            transforms=self.valid_transforms,\n                                            )\n        elif stage == 'test':\n            self.test_dataset = ImageDataset(image_names=valid_df.image.values, \n                                            labels=valid_df[self.CFG.labels].values, \n                                            image_dir=self.image_dir, \n                                            transforms=self.valid_transforms,\n                                            )\n        \n        \n    def train_dataloader(self):\n#         if self.trainer and self.trainer.current_epoch == 0:\n#             print('Set to none-augmentation transforms')\n#             self.train_dataset.transforms = self.no_aug_transforms\n#         else:\n#             print('Set to augmentation transforms')\n#             self.train_dataset.transforms = self.train_transforms\n        \n        train_loader = DataLoader(\n            self.train_dataset,\n            batch_size=self.CFG.batch_size,\n            num_workers=self.CFG.num_workers,\n            shuffle=True,\n            pin_memory=True,\n        )\n        return train_loader\n\n    def val_dataloader(self):        \n        valid_loader = DataLoader(\n            self.valid_dataset,\n            batch_size=self.CFG.batch_size,\n            num_workers=self.CFG.num_workers,\n            shuffle=False,\n            pin_memory=True,\n        )\n        return valid_loader\n\n    def test_dataloader(self):\n        return None","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Hàm chuyển label dạng text list sang label vector\ndef one_hot_encoded_df(dataset_df):\n    # copy dataframe\n    dataset_df_copy = dataset_df.copy()\n    unique_labels = dataset_df_copy.labels.unique()\n    new_column_names = list(set(' '.join(unique_labels).split()))\n    # initialize columns with zero\n    dataset_df_copy[new_column_names] = 0        \n    # one-hot-encoding using the column names\n    for labels in unique_labels:                \n        label_indices = dataset_df_copy[dataset_df_copy['labels'] == labels].index\n        splited_labels = labels.split()\n        dataset_df_copy.loc[label_indices, splited_labels] = 1\n    return dataset_df_copy","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load CSVs","metadata":{}},{"cell_type":"code","source":"dataset_df = pd.read_csv(CFG.train_csv_path)\nfolds_df = pd.read_csv(CFG.folds_csv_path)\n\ndf = one_hot_encoded_df(dataset_df)\ndf = folds_df.merge(df, on='image')\ndel dataset_df, folds_df\ndf.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Data transforms","metadata":{}},{"cell_type":"code","source":"# Wrapper để sử dụng CoarseDropout của albumentation cùng Transform mặc định của Pytorch\nclass AlbWrapper(object):\n    def __init__(self):\n        self.tf = A.CoarseDropout(max_height=int(360 * 0.08), max_width=int(360 * 0.08), max_holes=5, p=0.3)\n    def __call__(self, img):\n        return Image.fromarray(self.tf(image=np.array(img))['image'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train_aug = A.Compose([\n#         # A.RandomResizedCrop(360, 512, scale=(0.9, 1), p=0.5), \n#         A.Flip(p=0.5),\n#         A.OneOf([ \n#             A.ShiftScaleRotate(scale_limit=(-0.1, 0.00), rotate_limit=10),\n#             A.Perspective(scale=(0.05, 0.2)),\n#         ],p=0.3),\n#         A.OneOf([\n#             A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10),\n#             A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.2, 0.2)),     \n#         ], p=0.3),\n#         A.CoarseDropout(max_height=int(360 * 0.08), max_width=int(360 * 0.08), max_holes=5, p=0.3),\n# ])\n\ntrain_aug = tsfm.Compose([tsfm.RandomApply([tsfm.RandomPerspective(distortion_scale=0.2),], p=0.3),\n                        tsfm.RandomApply([tsfm.ColorJitter(0.2, 0.2, 0.2),tsfm.RandomAffine(degrees=10),], p=0.3),\n                        tsfm.RandomVerticalFlip(p=0.5),\n                        tsfm.RandomHorizontalFlip(p=0.5),\n                        AlbWrapper(),\n                     ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Plot một số ví dụ về Augmentation\n","metadata":{}},{"cell_type":"code","source":"def plot_sample(ds):\n    figure, axes = plt.subplots(3, 6, figsize=[20, 8])\n    for i, ax in enumerate(axes.flat):\n        image, _ = ds[i]\n        ax.imshow(image)\n        ax.axis('off')\n    plt.show()\n\nds = ImageDataset(image_names=df.image.values, \n                  labels=df[CFG.labels].values, \n                  image_dir=CFG.train_imgs_dir, \n                  transforms=None)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# No Aug\nplot_sample(ds)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train Aug\nds.transforms = train_aug\nplot_sample(ds)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del ds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Train và Valid Transform hoàn chỉnh","metadata":{}},{"cell_type":"code","source":"# train_transform = A.Compose([\n#       train_aug, \n#       A.Normalize(),\n#       ToTensorV2(),\n# ])\n\n# valid_transform = A.Compose([\n# #     A.Resize(height=CFG2.img_size, width=CFG2.img_size, p=1.0),\n#     A.Normalize(),\n#     ToTensorV2(),\n# ])\n\nDATASET_IMAGE_MEAN = (0.485, 0.456, 0.406)\nDATASET_IMAGE_STD = (0.229, 0.224, 0.225)\ntrain_transform = tsfm.Compose([train_aug,\n                                tsfm.ToTensor(),\n                                tsfm.Normalize(DATASET_IMAGE_MEAN, DATASET_IMAGE_STD), ])\n\nvalid_transform = tsfm.Compose([tsfm.ToTensor(),\n                                tsfm.Normalize(DATASET_IMAGE_MEAN, DATASET_IMAGE_STD), ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build Model","metadata":{}},{"cell_type":"code","source":"# f = -1\n# count = df[df.fold != f][CFG.labels].values.sum(axis=0)\n# pos_weight = ((len(df[df.fold != f].index) - count) / count)\n# pos_weight = torch.round(torch.FloatTensor(pos_weight))\n# # pos_weight = torch.clip(pos_weight, 0, 5)\n# print(pos_weight)\n# pos_weight = torch.log10(pos_weight) + 1\n# print(pos_weight)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Focal Loss\n\nImplement cùng class weight và pos_weight","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import Tensor\n\nclass FocalLoss(nn.modules.loss._Loss):\n    def __init__(self, gamma=2., alpha=1., weight = None, size_average=None, reduce=None, reduction = 'mean',\n                 pos_weight = None) -> None:\n        super(FocalLoss, self).__init__(size_average, reduce, reduction)\n        self.register_buffer('weight', weight)\n        self.register_buffer('pos_weight', pos_weight)\n        self.alpha = alpha\n        self.gamma = gamma\n        self.use_pw = True\n    \n    def reduce_loss(self, loss):\n        return loss.mean() if self.reduction == 'mean' else loss.sum() \\\n         if self.reduction == 'sum' else loss\n\n    def forward(self, input, target):\n        assert self.weight is None or isinstance(self.weight, Tensor)\n        assert self.pos_weight is None or isinstance(self.pos_weight, Tensor)\n        logpt = -F.binary_cross_entropy_with_logits(input, target,\n#                                                       self.weight,\n                                                      pos_weight=self.pos_weight if self.use_pw else None,\n                                                      reduction='none')\n        pt = torch.exp(logpt)\n\n        # compute the loss\n        focal_loss = -( self.alpha * (1-pt)**self.gamma ) * logpt\n        return self.reduce_loss(focal_loss * self.weight)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pytorch_lightning.callbacks import Callback\n\n# Callback điều chỉnh sử dụng pos weight và đổi learning rate\nclass RemovePosWeight(Callback):\n    def on_train_epoch_start(self, trainer, pl_module):\n        if pl_module.current_epoch == 1:\n            pass\n#             pl_module.criterion.use_pw = False\n#             print('PW off')\n#             trainer.optimizers[0].param_groups[0]['initial_lr'] = 5e-4\n#             trainer.optimizers[0].param_groups[0]['lr'] = 5e-4\n#             trainer.lr_schedulers[0]['scheduler'].base_lrs = [5e-4]\n#             print(trainer.lr_schedulers)\n#             print(trainer.optimizers[0])\n\n        elif pl_module.current_epoch == 2:\n#             trainer.lr_schedulers.clear()\n            trainer.optimizers[0].param_groups[0]['initial_lr'] = 1e-4\n            trainer.optimizers[0].param_groups[0]['lr'] = 1e-4\n            trainer.lr_schedulers[0]['scheduler'].base_lrs = [1e-4]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### F1 Metrics\n\nĐược tính dựa trên quan sát từ https://www.kaggle.com/buinyi/understanding-the-evaluation-metric-cv","metadata":{}},{"cell_type":"code","source":"# F1 score metric\nclass F1Score(Metric):\n    def __init__(self, threshold: float = 0.5, dist_sync_on_step=False):\n        super().__init__(dist_sync_on_step=dist_sync_on_step)\n        self.threshold = threshold\n        self.add_state(\"tp\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n        self.add_state(\"fp\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n        self.add_state(\"fn\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n\n    def update(self, preds: torch.Tensor, target: torch.Tensor, sigmoid=True):\n        assert preds.shape == target.shape\n        with torch.no_grad():\n            if sigmoid: preds = torch.sigmoid(preds)\n            preds = (preds > self.threshold).type(torch.long)\n\n            target_healthy = 1 - torch.clip(target.sum(dim=-1, keepdim=True), 0, 1)\n            pred_healthy = 1 - torch.clip(preds.sum(dim=-1, keepdim=True), 0, 1)\n            preds = torch.cat([preds, pred_healthy], -1)\n            target = torch.cat([target, target_healthy], -1)\n\n            tp = (preds*target).sum()\n            fp = preds.sum() - tp\n            fn = ((1 - preds)*target).sum()\n        \n        self.tp += tp.item()\n        self.fp += fp.item()\n        self.fn += fn.item()\n\n    def compute(self):\n        f1 = 2.0 * self.tp / (2.0 * self.tp + self.fn + self.fp)\n        return f1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Pytorch Lighting Module\n","metadata":{}},{"cell_type":"code","source":"# Sử dụng label smoothing\n# Dựa trên https://github.com/pytorch/pytorch/issues/7455\ndef smooth_target(target, smoothing=0.1):\n    with torch.no_grad():\n        true_dist = torch.abs(target - smoothing)\n    return true_dist","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Lit(pl.LightningModule):\n    def __init__(self, cfg):\n        super(Lit, self).__init__()\n        self.cfg = cfg\n        self.model = timm.create_model(cfg.model_name, pretrained=True, num_classes=cfg.num_classes)\n        self.criterion = FocalLoss(alpha=cfg.fl_alpha, gamma=cfg.fl_gamma, \n                                   weight=torch.tensor(CFG.cls_weight), \n                                   pos_weight=torch.tensor(CFG.pos_weight))\n        self.metric = F1Score()\n       \n    def forward(self, x):\n        return self.model(x)\n    \n    def configure_optimizers(self):\n        if self.cfg.use_sgd:\n            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.cfg.lr, momentum=0.9)\n        else:\n            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.cfg.lr)\n        \n        scheduler = {\n              'scheduler': torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer,\n                                                                      T_0=self.cfg.t_max,\n                                                                      eta_min=self.cfg.min_lr,\n                                                                      verbose=False),\n              'interval':'step',\n              'frequency': CFG.scheduler_freq,\n              'monitor': 'valid_loss',\n        }\n        return [self.optimizer], [scheduler]\n    \n    def training_step(self, batch, batch_idx):\n        images, targets = batch\n        logits = self.model(images)\n        loss = self.criterion(logits, smooth_target(targets))\n        score = self.metric(logits, targets)\n        \n        logs = {'train_loss': loss, 'lr': self.optimizer.param_groups[0]['lr']}\n        self.log_dict(logs, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n        self.log( 'train_f1', score, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        images, targets = batch\n        logits = self.model(images)\n        loss = self.criterion(logits, smooth_target(targets))\n        score = self.metric(logits, targets)\n        \n        logs = {'valid_loss': loss, 'valid_f1': score}\n        self.log_dict(logs, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        return loss","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"markdown","source":"Ở đây ta sẽ chỉ chạy thử một vài batches","metadata":{}},{"cell_type":"code","source":"# import wandb\n# wandb.login(key='KEY')\nfrom pytorch_lightning.loggers import WandbLogger","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Chuyển sang chế độ debug sẽ tắt logger và checkpoint callback\ndebug = True","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Khởi tạo data module\ndata_module = ImageDataModule(df=df[df.fold != CFG.reserve_fold],\n                              train_transforms=train_transform,\n                              valid_transforms=valid_transform,\n                              no_aug_transforms=valid_transform,\n                              image_dir= CFG.train_imgs_dir,\n                              fold_num=0,\n                              configurations=CFG)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train nhiều fold\nfor fold_idx in range(0, CFG.n_train_fold):\n    # Set valid fold \n    data_module.fold_num = fold_idx\n    data_module.setup('fit')\n    \n    # Logger và checkpoint\n    if not debug:\n        logger = WandbLogger(project='PP2021_0', \n                            name=f'{CFG.model_name}_f{fold_idx}_{CFG.version}',\n                            id=f'{CFG.model_name}_f{fold_idx}_{CFG.version}')\n        logger.log_hyperparams(CFG.__dict__)\n        checkpoint_callback = ModelCheckpoint(dirpath=os.path.join('ckpt', f'{CFG.model_name}_{CFG.version}'),\n                                              monitor='valid_f1',\n                                              save_top_k=1,\n                                              save_last=False,\n                                              save_weights_only=True,\n                                              filename=f'f{CFG.model_name}_' + '{epoch:02d}-{valid_f1:.4f}',\n                                              verbose=False,\n                                              mode='max')\n    else:\n        logger = checkpoint_callback = False\n#     early_stop_callback = EarlyStopping(monitor='valid_loss', min_delta=CFG.early_stop_delta, patience=3, mode='min')\n    \n    # Trainer\n    trainer = Trainer(max_epochs=CFG.num_epochs,\n                      gpus=0, # Do hết thời gian chạy kaggle rồi\n                      # tpu_cores=8,\n                      accumulate_grad_batches=CFG.accum_grad_batch,\n                      callbacks=[RemovePosWeight()],\n                      checkpoint_callback=checkpoint_callback,\n                      logger=logger,\n                      \n                      # Sẽ chỉ chạy 10 train batches và 10 val batches\n                      limit_train_batches = 10,\n                      limit_val_batches = 10,\n                      \n                      # precision=16,\n                      # resume_from_checkpoint='/content/drive/MyDrive/PP2021/ckpt/tf_efficientnet_b4_ns_f0_col_pos_weight_v5/last.ckpt',\n                      # reload_dataloaders_every_epoch=True,\n                      weights_summary='top')\n    # Model\n    model = Lit(CFG)\n    # for param in model.model.conv_stem.parameters():\n    #     param.requires_grad = False\n    # for param in model.model.blocks.parameters():\n    #     param.requires_grad = False\n   \n    # Fit\n    trainer.fit(model, data_module)\n    if not debug:\n        logger.close()\n        wandb.finish()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}