{"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-30T06:30:43.942803Z","iopub.execute_input":"2024-10-30T06:30:43.943421Z","iopub.status.idle":"2024-10-30T06:31:00.989079Z","shell.execute_reply.started":"2024-10-30T06:30:43.943386Z","shell.execute_reply":"2024-10-30T06:31:00.988019Z"},"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\nfrom torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as FSDP\nos.environ.pop('TPU_PROCESS_ADDRESSES')\n#os.environ.pop('CLOUD_TPU_TASK_ID')","metadata":{"execution":{"iopub.status.busy":"2024-10-30T06:31:00.990712Z","iopub.execute_input":"2024-10-30T06:31:00.990971Z","iopub.status.idle":"2024-10-30T06:31:39.003783Z","shell.execute_reply.started":"2024-10-30T06:31:00.990945Z","shell.execute_reply":"2024-10-30T06:31:39.002904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.login()","metadata":{"execution":{"iopub.status.busy":"2024-10-30T06:31:39.004888Z","iopub.execute_input":"2024-10-30T06:31:39.005185Z","iopub.status.idle":"2024-10-30T06:31:53.675983Z","shell.execute_reply.started":"2024-10-30T06:31:39.005153Z","shell.execute_reply":"2024-10-30T06:31:53.675184Z"},"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-30T06:31:53.677827Z","iopub.execute_input":"2024-10-30T06:31:53.678469Z","iopub.status.idle":"2024-10-30T06:31:53.685276Z","shell.execute_reply.started":"2024-10-30T06:31:53.678437Z","shell.execute_reply":"2024-10-30T06:31:53.684465Z"},"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-30T06:31:53.686279Z","iopub.execute_input":"2024-10-30T06:31:53.686607Z","iopub.status.idle":"2024-10-30T06:31:53.707497Z","shell.execute_reply.started":"2024-10-30T06:31:53.686578Z","shell.execute_reply":"2024-10-30T06:31:53.706799Z"},"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-30T06:31:53.708461Z","iopub.execute_input":"2024-10-30T06:31:53.708708Z","iopub.status.idle":"2024-10-30T06:31:53.738974Z","shell.execute_reply.started":"2024-10-30T06:31:53.708684Z","shell.execute_reply":"2024-10-30T06:31:53.738283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import lightning as L\nfrom lightning.fabric.strategies import FSDPStrategy\nfrom lightning.fabric import loggers as f_loggers\nfrom wandb.integration.lightning.fabric import WandbLogger\nseed_everything(2024)\n\n\ndef train(fabric):\n    data_module = DogBreedDataModule()\n    data_module.setup(stage='fit')\n    train_dataloader = data_module.train_dataloader()\n\n    model = LitSystem()\n    optimizer, scheduler = model.configure_optimizers()\n    fabric.print(f\"Start fablic setup\")\n    model, optimizer = fabric.setup(model, optimizer[0])\n    fabric.print(f\"Start fablic setup dataloader\")\n    train_dataloader = fabric.setup_dataloaders(train_dataloader)\n    fabric.print(f\"Start train loop\")\n    model.on_train_start()\n    num_epochs = 10\n    for epoch in range(num_epochs):\n        for i, batch in enumerate(train_dataloader):\n            fabric.print(f\"{epoch}/{num_epochs}| Steps: {i}\")\n            #inputs, labels = batch\n            optimizer.zero_grad()\n            #outputs = model(inputs)\n            #loss = torch.nn.functional.cross_entropy(outputs, labels)\n            loss = model.training_step(batch, i)\n            fabric.backward(loss)\n            optimizer.step()\n\ntb_logger = f_loggers.TensorBoardLogger(root_dir='tb_log/')\ncsv_logger = f_loggers.CSVLogger(root_dir='csv_log/')\nwandb_logger = WandbLogger(project='tpu-test')\n\nfabric = L.Fabric(accelerator=\"tpu\", devices=8, strategy=\"xla\", precision=\"bf16-true\", \n                  loggers=wandb_logger,\n                  #loggers=[tb_logger, csv_logger, wandb_logger]\n                 )\nfabric.launch(train)","metadata":{"execution":{"iopub.status.busy":"2024-10-30T06:38:16.333282Z","iopub.execute_input":"2024-10-30T06:38:16.333982Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2024-10-30T06:37:42.404915Z","iopub.execute_input":"2024-10-30T06:37:42.405605Z","iopub.status.idle":"2024-10-30T06:37:47.023235Z","shell.execute_reply.started":"2024-10-30T06:37:42.405566Z","shell.execute_reply":"2024-10-30T06:37:47.022598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}