{"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":[{"id":"477ec17c","cell_type":"code","source":"import os\nimport gc\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nimport timm\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt","metadata":{},"outputs":[],"execution_count":null},{"id":"9113f4dc","cell_type":"code","source":"# ==========================================\n# 1. CẤU HÌNH ĐƯỜNG DẪN \n# ==========================================\n# ĐƯỜNG DẪN DỮ LIỆU GỐC\nDATA_DIR = '/kaggle/input/competitions/malware-classification'\nLABELS_CSV = os.path.join(DATA_DIR, 'trainLabels.csv')\nSAMPLE_SUB_CSV = os.path.join(DATA_DIR, 'sampleSubmission.csv')\n\nTRAIN_BYTES_DIR = '/kaggle/input/datasets/trankimhuu/images-2-datasets-of-big2015/train/train/train_images_bytes' \nTRAIN_ASM_DIR = '/kaggle/input/datasets/trankimhuu/images-2-datasets-of-big2015/train/train/train_images_asm'\n\nTEST_BYTES_DIR = '/kaggle/input/datasets/trankimhuu/images-2-datasets-of-big2015/test/test/test_images_bytes'\nTEST_ASM_DIR = '/kaggle/input/datasets/trankimhuu/images-2-datasets-of-big2015/test/test/test_images_asm'\n\n# THƯ MỤC LƯU MODEL\nMODEL_SAVE_PATH = '/kaggle/working/best_dual_model.pth'\nOUTPUT_SUBMISSION = '/kaggle/working/submission.csv'\n\n# ==========================================\n# 2. HYPERPARAMETERS & 2-PHASE STRATEGY\n# ==========================================\nBATCH_SIZE = 32\nLEARNING_RATE = 1e-4\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Phân chia số Epoch cho từng giai đoạn\nEPOCHS_PHASE_1 = 5    # Đóng băng Backbone, chỉ train lớp phân loại\nEPOCHS_PHASE_2 = 45   # Mở băng toàn bộ, tinh chỉnh (Fine-tuning)\nTOTAL_EPOCHS = EPOCHS_PHASE_1 + EPOCHS_PHASE_2\n\nEARLY_STOP_PATIENCE = 7","metadata":{},"outputs":[],"execution_count":null},{"id":"3ebe20e8","cell_type":"code","source":"from sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\n\nclass DualBranchDataset(Dataset):\n    def __init__(self, df, bytes_dir, asm_dir, transform=None, is_test=False):\n        self.df = df\n        self.bytes_dir = bytes_dir\n        self.asm_dir = asm_dir\n        self.transform = transform\n        self.is_test = is_test\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        img_id = str(self.df.iloc[idx]['Id'])\n        \n        bytes_path = os.path.join(self.bytes_dir, f\"{img_id}.png\")\n        asm_path = os.path.join(self.asm_dir, f\"{img_id}.png\")\n        \n        try:\n            img_bytes = Image.open(bytes_path).convert('RGB')\n        except FileNotFoundError:\n            img_bytes = Image.new('RGB', (256, 256), color='black')\n            \n        try:\n            img_asm = Image.open(asm_path).convert('RGB')\n        except FileNotFoundError:\n            img_asm = Image.new('RGB', (256, 256), color='black')\n            \n        if self.transform:\n            img_bytes = self.transform(img_bytes)\n            img_asm = self.transform(img_asm)\n            \n        if self.is_test:\n            return img_bytes, img_asm, img_id\n        else:\n            label = int(self.df.iloc[idx]['Class']) - 1\n            return img_bytes, img_asm, label\n\n# Transform chung\nbase_transform = transforms.Compose([\n    transforms.Resize((256, 256)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n])\n\n# Chỉ đọc df_train_full, việc chia Fold sẽ thực hiện bên trong hàm huấn luyện\ndf_train_full = pd.read_csv(LABELS_CSV)","metadata":{},"outputs":[],"execution_count":null},{"id":"a7037270","cell_type":"code","source":"class DualBranchMalwareNet(nn.Module):\n    def __init__(self, num_classes=9):\n        super().__init__()\n        \n        self.branch_bytes = timm.create_model('swinv2_tiny_window8_256', pretrained=True, num_classes=0)\n        bytes_out_dim = self.branch_bytes.num_features \n        \n        self.branch_asm = timm.create_model('resnet18', pretrained=True, num_classes=0)\n        asm_out_dim = self.branch_asm.num_features \n        \n        self.classifier = nn.Sequential(\n            nn.Dropout(p=0.4),\n            nn.Linear(bytes_out_dim + asm_out_dim, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(p=0.4),\n            nn.Linear(512, num_classes)\n        )\n\n    def forward(self, img_bytes, img_asm):\n        feat_bytes = self.branch_bytes(img_bytes)\n        feat_asm = self.branch_asm(img_asm)\n        combined_features = torch.cat((feat_bytes, feat_asm), dim=1)\n        out = self.classifier(combined_features)\n        return out","metadata":{},"outputs":[],"execution_count":null},{"id":"0d7bd184","cell_type":"code","source":"def train_kfold():\n    # Khởi tạo Stratified K-Fold k=5\n    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n    fold_histories = []\n    \n    for fold, (train_idx, val_idx) in enumerate(skf.split(df_train_full, df_train_full['Class'])):\n        print(f\"\\n{'='*50}\")\n        print(f\"BẮT ĐẦU HUẤN LUYỆN FOLD {fold+1}/5\")\n        print(f\"{'='*50}\")\n        \n        df_train = df_train_full.iloc[train_idx]\n        df_val = df_train_full.iloc[val_idx]\n        \n        train_loader = DataLoader(DualBranchDataset(df_train, TRAIN_BYTES_DIR, TRAIN_ASM_DIR, base_transform), \n                                  batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n        val_loader = DataLoader(DualBranchDataset(df_val, TRAIN_BYTES_DIR, TRAIN_ASM_DIR, base_transform), \n                                batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n        \n        # Khởi tạo model sạch từ đầu cho mỗi fold\n        model = DualBranchMalwareNet(num_classes=9).to(DEVICE)\n        \n        # Áp dụng Label Smoothing chống Overfit\n        criterion = nn.CrossEntropyLoss(label_smoothing=0.1)\n        \n        best_val_loss = float('inf')\n        early_stop_counter = 0\n        history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}\n        fold_model_path = f'/kaggle/working/best_dual_model_fold{fold+1}.pth'\n        \n        for epoch in range(TOTAL_EPOCHS):\n            # [PHASE 1]\n            if epoch == 0:\n                print(f\"\\n[PHASE 1] FREEZE BACKBONE (Epoch 1 - {EPOCHS_PHASE_1}) - Fold {fold+1}\")\n                for param in model.branch_bytes.parameters(): param.requires_grad = False\n                for param in model.branch_asm.parameters(): param.requires_grad = False\n                optimizer = optim.AdamW(model.classifier.parameters(), lr=1e-3, weight_decay=1e-2)\n                \n            # [PHASE 2]\n            elif epoch == EPOCHS_PHASE_1:\n                print(f\"\\n[PHASE 2] UNFREEZE TOÀN MẠNG (Epoch {EPOCHS_PHASE_1 + 1} - {TOTAL_EPOCHS}) - Fold {fold+1}\")\n                for param in model.branch_bytes.parameters(): param.requires_grad = True\n                for param in model.branch_asm.parameters(): param.requires_grad = True\n                \n                optimizer = optim.AdamW([\n                    {'params': model.branch_bytes.parameters(), 'lr': 1e-5},\n                    {'params': model.branch_asm.parameters(), 'lr': 1e-5},\n                    {'params': model.classifier.parameters(), 'lr': 1e-4}\n                ], weight_decay=1e-2)\n\n            # ============ TRAIN ============\n            model.train()\n            train_loss, train_correct, train_total = 0.0, 0, 0\n            \n            for bytes_imgs, asm_imgs, labels in tqdm(train_loader, desc=f\"Fold {fold+1} Epoch {epoch+1} [Train]\", leave=False):\n                bytes_imgs, asm_imgs, labels = bytes_imgs.to(DEVICE), asm_imgs.to(DEVICE), labels.to(DEVICE)\n                \n                optimizer.zero_grad()\n                outputs = model(bytes_imgs, asm_imgs)\n                loss = criterion(outputs, labels)\n                loss.backward()\n                \n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n                optimizer.step()\n                \n                train_loss += loss.item()\n                _, preds = outputs.max(1)\n                train_correct += preds.eq(labels).sum().item()\n                train_total += labels.size(0)\n                \n            epoch_train_loss = train_loss / len(train_loader)\n            epoch_train_acc = train_correct / train_total\n            \n            # ============ VALIDATION ============\n            model.eval()\n            val_loss, val_correct, val_total = 0.0, 0, 0\n            \n            with torch.no_grad():\n                for bytes_imgs, asm_imgs, labels in tqdm(val_loader, desc=f\"Fold {fold+1} Epoch {epoch+1} [Val]\", leave=False):\n                    bytes_imgs, asm_imgs, labels = bytes_imgs.to(DEVICE), asm_imgs.to(DEVICE), labels.to(DEVICE)\n                    \n                    outputs = model(bytes_imgs, asm_imgs)\n                    loss = criterion(outputs, labels)\n                    \n                    val_loss += loss.item()\n                    _, preds = outputs.max(1)\n                    val_correct += preds.eq(labels).sum().item()\n                    val_total += labels.size(0)\n                    \n            epoch_val_loss = val_loss / len(val_loader)\n            epoch_val_acc = val_correct / val_total\n\n            history['train_loss'].append(epoch_train_loss)\n            history['train_acc'].append(epoch_train_acc)\n            history['val_loss'].append(epoch_val_loss)\n            history['val_acc'].append(epoch_val_acc)\n            \n            print(f\"Fold {fold+1} - Epoch {epoch+1:02d}/{TOTAL_EPOCHS} | \"\n                  f\"Train Loss: {epoch_train_loss:.4f} - Acc: {epoch_train_acc:.4f} | \"\n                  f\"Val Loss: {epoch_val_loss:.4f} - Acc: {epoch_val_acc:.4f}\")\n            \n            if epoch_val_loss < best_val_loss:\n                print(f\"LƯU MODEL FOLD {fold+1} (Val Loss giảm: {best_val_loss:.4f} -> {epoch_val_loss:.4f})\\n\")\n                best_val_loss = epoch_val_loss\n                torch.save(model.state_dict(), fold_model_path)\n                early_stop_counter = 0 \n            else:\n                early_stop_counter += 1\n                print(f\"Val Loss không giảm. Early Stop Counter: {early_stop_counter}/{EARLY_STOP_PATIENCE}\\n\")\n                if early_stop_counter >= EARLY_STOP_PATIENCE:\n                    print(f\"KÍCH HOẠT EARLY STOPPING! Dừng huấn luyện Fold {fold+1} sớm tại Epoch {epoch+1}.\\n\")\n                    break\n\n        fold_histories.append(history)\n        \n        # Dọn dẹp GPU VRAM chuẩn bị cho Fold tiếp theo\n        del model, optimizer, train_loader, val_loader\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    return fold_histories\n\ntraining_histories = train_kfold()","metadata":{},"outputs":[],"execution_count":null},{"id":"c0eb4116","cell_type":"code","source":"def inference_kfold_tta(n_folds=5, n_aug=4):\n    print(\"Bắt đầu Dự đoán: Ensemble 5 Folds + Test-Time Augmentation (TTA)...\")\n    \n    df_sub = pd.read_csv(SAMPLE_SUB_CSV)\n    \n    # Cấu hình 4 biến thể ảnh TTA: Gốc, Lật ngang, Lật dọc, Xoay 90 độ\n    augments = [\n        transforms.Lambda(lambda x: x), \n        transforms.RandomHorizontalFlip(p=1.0),\n        transforms.RandomVerticalFlip(p=1.0),\n        transforms.RandomRotation((90, 90)),\n    ]\n    \n    all_probs_ensemble = []\n    \n    for fold in range(1, n_folds + 1):\n        print(f\"\\nĐang nạp trọng số mô hình Fold {fold}...\")\n        fold_model_path = f'/kaggle/working/best_dual_model_fold{fold}.pth'\n        model = DualBranchMalwareNet(num_classes=9).to(DEVICE)\n        model.load_state_dict(torch.load(fold_model_path))\n        model.eval()\n        \n        fold_probs = []\n        \n        for aug_idx, aug in enumerate(augments[:n_aug]):\n            print(f\"  -> Chạy TTA version {aug_idx + 1}/{n_aug} trên Fold {fold}\")\n            \n            tta_transform = transforms.Compose([\n                transforms.Resize((256, 256)),\n                aug,\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n            ])\n            \n            test_loader = DataLoader(DualBranchDataset(df_sub, TEST_BYTES_DIR, TEST_ASM_DIR, tta_transform, is_test=True), \n                                     batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n            \n            tta_preds = []\n            all_ids = []\n            \n            with torch.no_grad():\n                for bytes_imgs, asm_imgs, img_ids in tqdm(test_loader, desc=f\"Testing Fold {fold} TTA {aug_idx+1}\", leave=False):\n                    bytes_imgs, asm_imgs = bytes_imgs.to(DEVICE), asm_imgs.to(DEVICE)\n                    \n                    outputs = model(bytes_imgs, asm_imgs)\n                    probs = F.softmax(outputs, dim=1)\n                    \n                    tta_preds.extend(probs.cpu().numpy())\n                    if fold == 1 and aug_idx == 0:\n                        all_ids.extend(img_ids)\n                        \n            if fold == 1 and aug_idx == 0:\n                final_ids = all_ids\n                \n            fold_probs.append(tta_preds)\n            \n        # Tính trung bình các phép TTA cho Fold hiện tại\n        fold_probs_avg = np.mean(fold_probs, axis=0)\n        all_probs_ensemble.append(fold_probs_avg)\n        \n    print(\"\\nĐang tính toán ma trận trung bình của 5 Folds (Ensemble)...\")\n    final_probs = np.mean(all_probs_ensemble, axis=0)\n    \n    col_names = [f'Prediction{i}' for i in range(1, 10)]\n    df_result = pd.DataFrame(final_probs, columns=col_names)\n    df_result.insert(0, 'Id', final_ids)\n    \n    df_result.to_csv(OUTPUT_SUBMISSION, index=False)\n    print(f\"Hoàn tất! File nộp bài Kaggle đã được lưu tại: {OUTPUT_SUBMISSION}\")\n\n# Chạy pipeline tạo file submission\ninference_kfold_tta()","metadata":{},"outputs":[],"execution_count":null}]}