{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":158833581,"sourceType":"kernelVersion"}],"dockerImageVersionId":30636,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torchaudio\nimport pyarrow.parquet as pq\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.transforms import transforms\nimport numpy as np\nfrom transformers import ViTModel\nimport pandas as pd\nimport torchvision.models as models\nimport os\nfrom tqdm import tqdm\nfrom torch.nn.functional import one_hot, softmax\nimport matplotlib.cm as cm\nfrom PIL import Image","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-01-13T16:56:11.023887Z","iopub.execute_input":"2024-01-13T16:56:11.024222Z","iopub.status.idle":"2024-01-13T16:56:11.030249Z","shell.execute_reply.started":"2024-01-13T16:56:11.024177Z","shell.execute_reply":"2024-01-13T16:56:11.029325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/hms-harmful-brain-activity-classification/train.csv')\ntest = pd.read_csv('/kaggle/input/hms-harmful-brain-activity-classification/test.csv')\nclasses = train['expert_consensus'].unique()\nmapping = {\n    c:i for i, c in enumerate(classes)\n}\nnum_classes = classes.shape[0]","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:11.031273Z","iopub.execute_input":"2024-01-13T16:56:11.031515Z","iopub.status.idle":"2024-01-13T16:56:11.217146Z","shell.execute_reply.started":"2024-01-13T16:56:11.031493Z","shell.execute_reply":"2024-01-13T16:56:11.216086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cmap = cm.get_cmap(\"viridis\")","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:11.219594Z","iopub.execute_input":"2024-01-13T16:56:11.219877Z","iopub.status.idle":"2024-01-13T16:56:11.224559Z","shell.execute_reply.started":"2024-01-13T16:56:11.219853Z","shell.execute_reply":"2024-01-13T16:56:11.223655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SpectrogramDataset(Dataset):\n    def __init__(self, data_folder, transform=None):\n        self.data_folder = data_folder\n        self.file_paths = [os.path.join(data_folder, f) for f in os.listdir(data_folder)]\n        self.transform = transform\n#         self.remove_empty()\n    \n    def __len__(self):\n        return len(self.file_paths)\n\n    def __getitem__(self, idx):\n        file_path = self.file_paths[idx]\n        label = os.path.basename(file_path).split(\"_\")[1].split(\".\")[0]  # Extract label from filename\n        spectrogram = pd.read_parquet(file_path).drop('time',axis=1).values  # Load parquet file\n        spectrogram = Image.fromarray((cmap(spectrogram) * 255).astype(np.uint8))\n        if self.transform:\n            spectrogram = self.transform(spectrogram)[:3, :, :]\n            \n        return spectrogram, mapping[label]","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:11.225947Z","iopub.execute_input":"2024-01-13T16:56:11.226394Z","iopub.status.idle":"2024-01-13T16:56:11.239042Z","shell.execute_reply.started":"2024-01-13T16:56:11.226359Z","shell.execute_reply":"2024-01-13T16:56:11.238243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_folder = \"/kaggle/input/hms-data-prepare-separate-spectogram/separate_spectogram/\"\n\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),  # Resize to 224x224\n#     transforms.CenterCrop(224),  # Center crop to maintain aspect ratio\n    transforms.ToTensor(),  # Convert to PyTorch tensor\n#     transforms.Normalize(mean=[0.485, 0.456, 0.406, 0.406], std=[0.229, 0.224, 0.225, 0.225])  # Normalize (optional)\n])\n\ndataset = SpectrogramDataset(data_folder, transform=transform)\ntrain_length = int(0.8 * len(dataset))\nval_length = len(dataset) - train_length\n\nprint(f'train size: {train_length} \\nval size: {val_length}')\n\ntrain_set, val_set = torch.utils.data.random_split(dataset, [train_length, val_length])\n\ntrain_loader = DataLoader(train_set, batch_size=32, shuffle=True)\nval_loader = DataLoader(val_set, batch_size=32)","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:26.259565Z","iopub.execute_input":"2024-01-13T16:56:26.259968Z","iopub.status.idle":"2024-01-13T16:56:26.319716Z","shell.execute_reply.started":"2024-01-13T16:56:26.259935Z","shell.execute_reply":"2024-01-13T16:56:26.318795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ViTClassifier(torch.nn.Module):\n    def __init__(self, num_classes=1000):\n        super().__init__()\n        self.vit = ViTModel.from_pretrained(\"google/vit-base-patch16-224-in21k\")\n        self.classifier = torch.nn.Linear(self.vit.config.hidden_size, num_classes)\n\n    def forward(self, images):\n        output = self.vit(images)\n        output = self.classifier(output.last_hidden_state[:, 0]) \n        output = softmax(output, dim = 1)\n        return output","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:26.321495Z","iopub.execute_input":"2024-01-13T16:56:26.321799Z","iopub.status.idle":"2024-01-13T16:56:26.327848Z","shell.execute_reply.started":"2024-01-13T16:56:26.321774Z","shell.execute_reply":"2024-01-13T16:56:26.326889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = 'cuda'\nmodel = ViTClassifier(num_classes).to(device)","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:26.329018Z","iopub.execute_input":"2024-01-13T16:56:26.329281Z","iopub.status.idle":"2024-01-13T16:56:27.023497Z","shell.execute_reply.started":"2024-01-13T16:56:26.329258Z","shell.execute_reply":"2024-01-13T16:56:27.022725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_fn = nn.KLDivLoss(reduction=\"batchmean\")\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)","metadata":{"execution":{"iopub.status.busy":"2024-01-13T16:56:27.025415Z","iopub.execute_input":"2024-01-13T16:56:27.025725Z","iopub.status.idle":"2024-01-13T16:56:27.031214Z","shell.execute_reply.started":"2024-01-13T16:56:27.025697Z","shell.execute_reply":"2024-01-13T16:56:27.030216Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs = 10\n\nfor epoch in range(num_epochs):\n    pbar = tqdm(train_loader)\n    for spectrograms, labels in pbar:\n        labels_onehot = one_hot(labels, num_classes=num_classes).float().to(device)\n        spectrograms = spectrograms.to(device)\n        optimizer.zero_grad()\n        outputs = model(spectrograms)\n        loss = loss_fn(outputs.log(), labels_onehot)\n        loss.backward()\n        optimizer.step()\n        pbar.set_description(f'train loss {loss}')\n\n    # Evaluate on validation set\n    with torch.no_grad():\n        correct = 0\n        total = 0\n        val_loss = 0\n        pbar = tqdm(val_loader)\n        for spectrograms, labels in pbar:\n            labels = labels.to(device)\n            labels_onehot = one_hot(labels, num_classes=num_classes).float()\n            spectrograms = spectrograms.to(device)\n            outputs = model(spectrograms)\n            _, predicted = torch.max(outputs.data, 1)\n            total += labels.size(0)\n            correct += (predicted == labels).sum().item()\n            loss = loss_fn(outputs.log(), labels_onehot).item()\n            val_loss += loss\n            pbar.set_description(f'val loss {loss}')\n\n        accuracy = 100 * correct / total\n        kl_divergence = val_loss / len(val_loader)\n\n        print(\"Validation Accuracy:\", accuracy)\n        print(\"Validation KL Divergence:\", kl_divergence)\n\n\n# Save the trained model\ntorch.save(model.state_dict(), \"trained_model.pt\")","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-01-13T16:56:27.032567Z","iopub.execute_input":"2024-01-13T16:56:27.032910Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}