{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install py7zr -q\nimport os\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ==================== 1. 路径配置 ====================\n# 数据路径（根据诊断结果）\ninput_dir = '/kaggle/input/competitions/cifar-10'\noutput_dir = '/kaggle/working'\nos.makedirs(output_dir, exist_ok=True)\n\n# 解压后的图片目录\ntrain_dir = os.path.join(output_dir, 'train')\ntest_dir = os.path.join(output_dir, 'test')\n\n# CSV文件路径（直接在input_dir下）\ntrain_labels_path = os.path.join(input_dir, 'trainLabels.csv')\nsample_submission_path = os.path.join(input_dir, 'sampleSubmission.csv')\n\nprint(f\"数据目录: {input_dir}\")\nprint(f\"训练标签: {os.path.exists(train_labels_path)}\")\nprint(f\"提交样例: {os.path.exists(sample_submission_path)}\")\n\n# ==================== 2. 解压数据（重要！）====================\nprint(\"\\n正在解压训练图片...\")\nif not os.path.exists(train_dir):\n    import py7zr\n    with py7zr.SevenZipFile(os.path.join(input_dir, 'train.7z'), mode='r') as z:\n        z.extractall(output_dir)\n    print(f\"训练图片解压完成: {train_dir}\")\n\nprint(\"正在解压测试图片...\")\nif not os.path.exists(test_dir):\n    import py7zr\n    with py7zr.SevenZipFile(os.path.join(input_dir, 'test.7z'), mode='r') as z:\n        z.extractall(output_dir)\n    print(f\"测试图片解压完成: {test_dir}\")\n\n# 验证解压结果\nprint(f\"\\n训练图片数量: {len(os.listdir(train_dir)) if os.path.exists(train_dir) else 0}\")\nprint(f\"测试图片数量: {len(os.listdir(test_dir)) if os.path.exists(test_dir) else 0}\")\n\n# ==================== 3. 设备配置 ====================\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'\\nUsing device: {device}')\nif torch.cuda.is_available():\n    print(f'GPU: {torch.cuda.get_device_name(0)}')\n\n# ==================== 4. 读取标签 ====================\ndf_labels = pd.read_csv(train_labels_path)\nprint(f'\\n训练标签数: {len(df_labels)}')\nprint(df_labels.head())\n\nclasses = ['airplane', 'automobile', 'bird', 'cat', 'deer', \n           'dog', 'frog', 'horse', 'ship', 'truck']\nclass_to_idx = {cls: i for i, cls in enumerate(classes)}\nidx_to_class = {i: cls for cls, i in class_to_idx.items()}\nlabel_dict = dict(zip(df_labels['id'].astype(str), df_labels['label']))\n\n# ==================== 5. 数据集类 ====================\nclass CIFAR10Dataset(Dataset):\n    def __init__(self, img_dir, labels=None, transform=None, is_test=False):\n        self.img_dir = img_dir\n        self.labels = labels\n        self.transform = transform\n        self.is_test = is_test\n        self.images = sorted([f for f in os.listdir(img_dir) if f.endswith('.png')])\n        \n    def __len__(self):\n        return len(self.images)\n    \n    def __getitem__(self, idx):\n        img_name = self.images[idx]\n        img_path = os.path.join(self.img_dir, img_name)\n        image = Image.open(img_path).convert('RGB')\n        \n        if self.transform:\n            image = self.transform(image)\n        \n        if self.is_test:\n            img_id = img_name.split('.')[0]\n            return image, img_id\n        else:\n            img_id = img_name.split('.')[0]\n            label_idx = class_to_idx[self.labels[img_id]]\n            return image, label_idx\n\n# ==================== 6. 数据增强 ====================\ntrain_transform = transforms.Compose([\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomCrop(32, padding=4),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], \n                        std=[0.2023, 0.1994, 0.2010])\n])\n\ntest_transform = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], \n                        std=[0.2023, 0.1994, 0.2010])\n])\n\n# ==================== 7. 划分验证集 ====================\nfrom sklearn.model_selection import train_test_split\n\nall_images = [f for f in os.listdir(train_dir) if f.endswith('.png')]\nall_labels = [label_dict[f.split('.')[0]] for f in all_images]\ntrain_images, val_images = train_test_split(\n    all_images, test_size=0.1, random_state=42, stratify=all_labels\n)\n\nclass FilteredDataset(Dataset):\n    def __init__(self, img_dir, image_list, labels, transform=None):\n        self.img_dir = img_dir\n        self.image_list = image_list\n        self.labels = labels\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.image_list)\n    \n    def __getitem__(self, idx):\n        img_name = self.image_list[idx]\n        img_path = os.path.join(self.img_dir, img_name)\n        image = Image.open(img_path).convert('RGB')\n        if self.transform:\n            image = self.transform(image)\n        img_id = img_name.split('.')[0]\n        label_idx = class_to_idx[self.labels[img_id]]\n        return image, label_idx\n\ntrain_dataset = FilteredDataset(train_dir, train_images, label_dict, transform=train_transform)\nval_dataset = FilteredDataset(train_dir, val_images, label_dict, transform=test_transform)\n\nprint(f\"训练集: {len(train_dataset)} 张\")\nprint(f\"验证集: {len(val_dataset)} 张\")\n\n# ==================== 8. 模型定义 (ResNet18 for CIFAR-10) ====================\nclass BasicBlock(nn.Module):\n    expansion = 1\n    def __init__(self, in_planes, planes, stride=1):\n        super(BasicBlock, self).__init__()\n        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.shortcut = nn.Sequential()\n        if stride != 1 or in_planes != self.expansion * planes:\n            self.shortcut = nn.Sequential(\n                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),\n                nn.BatchNorm2d(self.expansion * planes)\n            )\n    def forward(self, x):\n        out = torch.relu(self.bn1(self.conv1(x)))\n        out = self.bn2(self.conv2(out))\n        out += self.shortcut(x)\n        out = torch.relu(out)\n        return out\n\nclass ResNet(nn.Module):\n    def __init__(self, block, num_blocks, num_classes=10):\n        super(ResNet, self).__init__()\n        self.in_planes = 64\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(64)\n        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)\n        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)\n        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)\n        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)\n        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n        self.fc = nn.Linear(512 * block.expansion, num_classes)\n    def _make_layer(self, block, planes, num_blocks, stride):\n        strides = [stride] + [1] * (num_blocks - 1)\n        layers = []\n        for stride in strides:\n            layers.append(block(self.in_planes, planes, stride))\n            self.in_planes = planes * block.expansion\n        return nn.Sequential(*layers)\n    def forward(self, x):\n        out = torch.relu(self.bn1(self.conv1(x)))\n        out = self.layer1(out)\n        out = self.layer2(out)\n        out = self.layer3(out)\n        out = self.layer4(out)\n        out = self.avgpool(out)\n        out = out.view(out.size(0), -1)\n        out = self.fc(out)\n        return out\n\ndef ResNet18():\n    return ResNet(BasicBlock, [2, 2, 2, 2])\n\nmodel = ResNet18().to(device)\nprint(f'\\n模型参数量: {sum(p.numel() for p in model.parameters()):,}')\n\n# ==================== 9. 训练配置 ====================\nbatch_size = 128\nlearning_rate = 0.1\nnum_epochs = 50\n\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)\n\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)\n\n# ==================== 10. 训练函数 ====================\ndef train_epoch(model, loader, criterion, optimizer):\n    model.train()\n    total_loss, correct, total = 0, 0, 0\n    for images, labels in tqdm(loader, desc='训练'):\n        images, labels = images.to(device), labels.to(device)\n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n        _, predicted = outputs.max(1)\n        total += labels.size(0)\n        correct += predicted.eq(labels).sum().item()\n    return total_loss / len(loader), 100. * correct / total\n\ndef validate(model, loader, criterion):\n    model.eval()\n    total_loss, correct, total = 0, 0, 0\n    with torch.no_grad():\n        for images, labels in tqdm(loader, desc='验证'):\n            images, labels = images.to(device), labels.to(device)\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n            total_loss += loss.item()\n            _, predicted = outputs.max(1)\n            total += labels.size(0)\n            correct += predicted.eq(labels).sum().item()\n    return total_loss / len(loader), 100. * correct / total\n\n# ==================== 11. 开始训练 ====================\nprint('\\n' + '='*50)\nprint('开始训练')\nprint('='*50)\n\nbest_acc = 0\nfor epoch in range(num_epochs):\n    print(f'\\nEpoch {epoch+1}/{num_epochs}, LR: {scheduler.get_last_lr()[0]:.4f}')\n    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer)\n    val_loss, val_acc = validate(model, val_loader, criterion)\n    scheduler.step()\n    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')\n    print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')\n    if val_acc > best_acc:\n        best_acc = val_acc\n        torch.save(model.state_dict(), os.path.join(output_dir, 'best_model.pth'))\n        print(f'✓ 保存模型，验证准确率: {best_acc:.2f}%')\n\nprint(f'\\n训练完成！最佳验证准确率: {best_acc:.2f}%')\n\n# ==================== 12. 预测测试集 ====================\nprint('\\n加载最佳模型...')\nmodel.load_state_dict(torch.load(os.path.join(output_dir, 'best_model.pth')))\nmodel.eval()\n\ntest_dataset = CIFAR10Dataset(test_dir, labels=None, transform=test_transform, is_test=True)\ntest_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=2)\nprint(f'测试集图片数: {len(test_dataset)}')\n\nall_ids, all_preds = [], []\nwith torch.no_grad():\n    for images, img_ids in tqdm(test_loader, desc='预测'):\n        images = images.to(device)\n        outputs = model(images)\n        _, predicted = outputs.max(1)\n        all_ids.extend(img_ids)\n        all_preds.extend(predicted.cpu().numpy())\n\npredicted_labels = [idx_to_class[p] for p in all_preds]\n\n# ==================== 13. 生成提交文件 ====================\nsubmission_df = pd.DataFrame({'id': all_ids, 'label': predicted_labels})\nsubmission_df = submission_df.sort_values('id').reset_index(drop=True)\nsubmission_path = os.path.join(output_dir, 'submission.csv')\nsubmission_df.to_csv(submission_path, index=False)\n\nprint(f'\\n提交文件已保存: {submission_path}')\nprint(f'文件大小: {os.path.getsize(submission_path) / 1024 / 1024:.2f} MB')\nprint('\\n前5行预览:')\nprint(submission_df.head())\nprint('\\n后5行预览:')\nprint(submission_df.tail())\nprint('\\n完成！点击右侧Output面板下载 submission.csv 并提交到Kaggle')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-06-16T03:51:12.188326Z","iopub.execute_input":"2026-06-16T03:51:12.188721Z","iopub.status.idle":"2026-06-16T04:28:28.454772Z","shell.execute_reply.started":"2026-06-16T03:51:12.188694Z","shell.execute_reply":"2026-06-16T04:28:28.453746Z"}},"outputs":[],"execution_count":null}]}