{"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-30T09:08:33.375815Z","iopub.execute_input":"2024-10-30T09:08:33.376814Z","iopub.status.idle":"2024-10-30T09:08:51.074460Z","shell.execute_reply.started":"2024-10-30T09:08:33.376775Z","shell.execute_reply":"2024-10-30T09:08:51.073563Z"},"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-30T09:08:51.076165Z","iopub.execute_input":"2024-10-30T09:08:51.076466Z","iopub.status.idle":"2024-10-30T09:09:29.240264Z","shell.execute_reply.started":"2024-10-30T09:08:51.076433Z","shell.execute_reply":"2024-10-30T09:09:29.239524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.login()","metadata":{"execution":{"iopub.status.busy":"2024-10-30T09:09:29.241189Z","iopub.execute_input":"2024-10-30T09:09:29.241434Z","iopub.status.idle":"2024-10-30T09:09:33.381924Z","shell.execute_reply.started":"2024-10-30T09:09:29.241409Z","shell.execute_reply":"2024-10-30T09:09:33.381108Z"},"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-30T09:09:33.383666Z","iopub.execute_input":"2024-10-30T09:09:33.384205Z","iopub.status.idle":"2024-10-30T09:09:33.390936Z","shell.execute_reply.started":"2024-10-30T09:09:33.384175Z","shell.execute_reply":"2024-10-30T09:09:33.390115Z"},"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-30T09:09:33.391816Z","iopub.execute_input":"2024-10-30T09:09:33.392051Z","iopub.status.idle":"2024-10-30T09:09:33.407781Z","shell.execute_reply.started":"2024-10-30T09:09:33.392026Z","shell.execute_reply":"2024-10-30T09:09:33.407036Z"},"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-30T09:09:33.408633Z","iopub.execute_input":"2024-10-30T09:09:33.408874Z","iopub.status.idle":"2024-10-30T09:09:33.422513Z","shell.execute_reply.started":"2024-10-30T09:09:33.408847Z","shell.execute_reply":"2024-10-30T09:09:33.421780Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import lightning as L\nfrom lightning.fabric.strategies import XLAFSDPStrategy\nfrom torch_xla.distributed.fsdp.wrap import size_based_auto_wrap_policy\nimport functools\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    #\n    fabric.print(f\"Start fablic setup\")\n    model = fabric.setup_module(model)\n    optimizer = torch.optim.Adam(list(model.model.parameters()), lr=0.001)\n    optimizer = fabric.setup_optimizers(optimizer)\n    \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\n#auto_wrap_policy = functools.partial(size_based_auto_wrap_policy,min_num_params=2000)\nauto_wrap_policy = functools.partial(size_based_auto_wrap_policy,min_num_params=0)\nstrategy = XLAFSDPStrategy(auto_wrap_policy=auto_wrap_policy,_debug_print=True,mark_step_on_finalization=True,flatten_parameters=True)\n\nfabric = L.Fabric(accelerator=\"tpu\", devices=8, strategy=strategy, 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-30T09:13:07.715055Z","iopub.execute_input":"2024-10-30T09:13:07.716186Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2024-10-30T09:12:29.050384Z","iopub.execute_input":"2024-10-30T09:12:29.050810Z","iopub.status.idle":"2024-10-30T09:12:31.831890Z","shell.execute_reply.started":"2024-10-30T09:12:29.050767Z","shell.execute_reply":"2024-10-30T09:12:31.831209Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}