{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.15","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":7327,"databundleVersionId":861871,"sourceType":"competition"}],"dockerImageVersionId":30788,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#! pip install --upgrade pip\n#! pip install cloud-tpu-client https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.13-cp38-cp38m-linux_x86_64.whl\n#! pip install lightning\n#! pip install pytorch-lightning\n#! pip install omegaconf\n!pip install wandb\n!pip install lightning","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-30T02:00:04.728870Z","iopub.execute_input":"2024-10-30T02:00:04.729411Z","iopub.status.idle":"2024-10-30T02:00:22.553316Z","shell.execute_reply.started":"2024-10-30T02:00:04.729373Z","shell.execute_reply":"2024-10-30T02:00:22.552456Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport numpy as np\nimport pandas as pd\n\nimport pytorch_lightning as pl\nfrom torchmetrics import Accuracy\nfrom pytorch_lightning import loggers\nfrom pytorch_lightning import seed_everything\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import StratifiedKFold\nimport albumentations as A\nimport timm\nimport os\nimport wandb\nos.environ.pop('TPU_PROCESS_ADDRESSES')\n#os.environ.pop('CLOUD_TPU_TASK_ID')","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:00:22.555023Z","iopub.execute_input":"2024-10-30T02:00:22.555299Z","iopub.status.idle":"2024-10-30T02:01:02.038344Z","shell.execute_reply.started":"2024-10-30T02:00:22.555272Z","shell.execute_reply":"2024-10-30T02:01:02.037597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.login()","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:01:02.039265Z","iopub.execute_input":"2024-10-30T02:01:02.039517Z","iopub.status.idle":"2024-10-30T02:01:24.823998Z","shell.execute_reply.started":"2024-10-30T02:01:02.039491Z","shell.execute_reply":"2024-10-30T02:01:24.823244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DogBreedDataset(Dataset):\n    def __init__(self, dataframe, data_dir, transform=None, test=False):\n        self.data = dataframe.reset_index(drop=True)\n        self.data_dir = data_dir\n        self.transform = transform\n        self.test = test\n        cv2.setNumThreads(0)\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        img_name = os.path.join(self.data_dir, self.data.loc[idx, \"id\"] + \".\" + \"jpg\")\n        image = cv2.imread(img_name)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        image = self.transform(image=image)\n        image = torch.from_numpy(image[\"image\"].transpose(2, 0, 1))\n        if not self.test:\n            label = self.data.loc[idx, \"label\"]\n            return image, label\n        else:\n            return image","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:01:24.825050Z","iopub.execute_input":"2024-10-30T02:01:24.825571Z","iopub.status.idle":"2024-10-30T02:01:24.831678Z","shell.execute_reply.started":"2024-10-30T02:01:24.825542Z","shell.execute_reply":"2024-10-30T02:01:24.831049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DogBreedDataModule(pl.LightningDataModule):\n\n    def __init__(self):\n        super().__init__()\n        \n\n    # OPTIONAL, called only on 1 GPU/machine(for download or tokenize)\n    def prepare_data(self):\n        pass\n\n    # OPTIONAL, called for every GPU/machine\n    def setup(self, stage=None):\n        if stage == 'fit':\n            df = pd.read_csv('../input/dog-breed-identification/labels.csv')\n            \n            df['label'] = df.groupby(['breed']).ngroup()\n            self.breed2id_dict = df[['breed','label']].sort_values(by='label').set_index('breed').to_dict()['label'] \n            self.id2breed_dict = df[['breed','label']].sort_values(by='label').set_index('label').to_dict()['breed']  \n\n            kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=2021)\n            for fold, (train_index, val_index) in enumerate(kf.split(df.values, df[\"breed\"])):\n                df.loc[val_index, \"fold\"] = int(fold)\n            df[\"fold\"] = df[\"fold\"].astype(int)\n\n            train_df = df[df[\"fold\"] != 0]\n            valid_df = df[df[\"fold\"] == 0]\n            train_df = pd.concat([train_df]*10)\n            train_transform = A.Compose([\n                        A.Resize(height=448, width=448, interpolation=1, always_apply=False, p=1.0),\n                        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0),\n                        ])\n\n            valid_transform = A.Compose([\n                        A.Resize(height=448, width=448, interpolation=1, always_apply=False, p=1.0),\n                        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0)\n                        ])\n\n            self.train_dataset = DogBreedDataset(train_df, '../input/dog-breed-identification/train', transform=train_transform)\n            self.valid_dataset = DogBreedDataset(valid_df, '../input/dog-breed-identification/train', transform=valid_transform)\n\n    def train_dataloader(self):\n        return DataLoader(self.train_dataset, batch_size=2, num_workers=4, shuffle=True, pin_memory=True, drop_last=True)\n\n    def val_dataloader(self):\n        return DataLoader(self.valid_dataset, batch_size=2, num_workers=4, shuffle=False, pin_memory=True, drop_last=True)","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:01:24.833218Z","iopub.execute_input":"2024-10-30T02:01:24.833516Z","iopub.status.idle":"2024-10-30T02:01:24.849137Z","shell.execute_reply.started":"2024-10-30T02:01:24.833488Z","shell.execute_reply":"2024-10-30T02:01:24.848519Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LitSystem(pl.LightningModule):\n    def __init__(self):\n        super().__init__()\n        #self.conf = conf\n        #self.save_hyperparameters(conf)\n        #self.model = timm.create_model(model_name=\"tf_efficientnet_b7\", num_classes=120, pretrained=False, in_chans=3)\n        self.model = timm.create_model(model_name=\"eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\", num_classes=120, pretrained=False, in_chans=3)\n        self.criteria = torch.nn.CrossEntropyLoss()\n        self.accuracy = Accuracy(task=\"multiclass\", num_classes=120)\n\n    def forward(self, x):\n        # use forward for inference/predictions\n        return self.model(x)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)\n\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)\n        \n        return [optimizer], [scheduler]\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self.model(x)\n        loss = self.criteria(y_hat, y)\n        \n        self.log('train_loss', loss, on_epoch=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self.model(x)\n        loss = self.criteria(y_hat, y)\n        self.log('val_loss', loss, on_epoch=True)","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:01:24.849931Z","iopub.execute_input":"2024-10-30T02:01:24.850188Z","iopub.status.idle":"2024-10-30T02:01:24.865100Z","shell.execute_reply.started":"2024-10-30T02:01:24.850164Z","shell.execute_reply":"2024-10-30T02:01:24.864404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_everything(2024)\n\ntb_logger = loggers.TensorBoardLogger(save_dir='tb_log/')\ncsv_logger = loggers.CSVLogger(save_dir='csv_log/')\nwandb_logger = loggers.WandbLogger(project='tpu-test')\nlr_monitor = LearningRateMonitor(logging_interval='step')\ncheckpoint_callback = ModelCheckpoint(dirpath='ckpt', monitor='val_loss', \n                                      save_last=True, save_top_k=5, mode='min', \n                                      save_weights_only=True, filename='{epoch}-{avg_val_loss:.2f}')\n\ndata_module = DogBreedDataModule()\n\nlit_model = LitSystem()\n\ntrainer = Trainer(\n    logger=[tb_logger, csv_logger, wandb_logger],\n    callbacks=[lr_monitor, checkpoint_callback],\n    max_epochs=10,\n    devices=\"auto\",\n    precision=\"bf16-true\"\n        )\n\ntrainer.fit(lit_model, data_module)","metadata":{"execution":{"iopub.status.busy":"2024-10-30T02:01:24.866004Z","iopub.execute_input":"2024-10-30T02:01:24.866266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}