{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84616,"databundleVersionId":9494719,"sourceType":"competition"},{"sourceId":9357079,"sourceType":"datasetVersion","datasetId":5672697},{"sourceId":109867,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":92016,"modelId":116221}],"dockerImageVersionId":30761,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision\nimport torchvision.transforms as transforms\nfrom tqdm import tqdm\nfrom timm import create_model  # Import Vision Transformer from timm library\n# Set device\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n# 运行该session后，请直接运行最后的eval代码！\n# Define transformations for CIFAR-10 dataset\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),  # Resize to 224x224 for Vision Transformer input\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize images\n])","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-09-22T03:41:27.343192Z","iopub.execute_input":"2024-09-22T03:41:27.343594Z","iopub.status.idle":"2024-09-22T03:41:27.350885Z","shell.execute_reply.started":"2024-09-22T03:41:27.343557Z","shell.execute_reply":"2024-09-22T03:41:27.349870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\n# Load CIFAR-10 dataset\ntrainset = torchvision.datasets.CIFAR10(root='/kaggle/input/2024-zju', train=True, download=True, transform=transform)\ntrainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)\n\n# Load Vision Transformer model from timm\nnet = create_model('vit_base_patch16_224', pretrained=True, num_classes=10)  # CIFAR-10 has 10 classes\nnet = net.to(device)\n\n# Loss function and optimizer\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.Adam(net.parameters(), lr=1e-4)\n\n# Train the model\ndef train_model(num_epochs=1):\n    net.train()\n    for epoch in range(num_epochs):\n        running_loss = 0.0\n        for i, (inputs, labels) in tqdm(enumerate(trainloader)):\n            inputs, labels = inputs.to(device), labels.to(device)\n            \n            optimizer.zero_grad()\n            outputs = net(inputs)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n\n            running_loss += loss.item()\n            if i % 10 == 1:    # Print every 100 mini-batches\n                print(f'[Epoch {epoch+1}, Batch {i+1}] Loss: {running_loss / 100:.3f}')\n                running_loss = 0.0\n            if i > 100:\n                break\n\n    print('Finished Training')\n    # Save the trained model weights\n    torch.save(net.state_dict(), 'cifar10_vit.pth')\n    print('Model weights saved!')\n\n# train_model()\n# I pretrained the model with 1 H100 on my on cluster\n# the weight has been created and uploaded on Kaggle","metadata":{"execution":{"iopub.status.busy":"2024-09-08T15:53:42.132490Z","iopub.execute_input":"2024-09-08T15:53:42.132955Z","iopub.status.idle":"2024-09-08T15:54:03.520931Z","shell.execute_reply.started":"2024-09-08T15:53:42.132909Z","shell.execute_reply":"2024-09-08T15:54:03.519230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pickle\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom PIL import Image\n\nclass CIFAR10TestDataset(Dataset):\n    def __init__(self, data_path, transform=None):\n        # 加载 test_batch 数据\n        with open(data_path, 'rb') as f:\n            data_dict = pickle.load(f, encoding='bytes')\n        \n        self.data = data_dict[b'data']  # 图像数据\n        self.labels = data_dict[b'labels']  # 标签\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.labels)\n\n    def __getitem__(self, idx):\n        # 获取单个样本的数据和标签\n        img_data = self.data[idx]\n        img_data = img_data.reshape(3, 32, 32)  # CIFAR-10 图像大小为 32x32，RGB 通道\n        img_data = img_data.transpose(1, 2, 0)  # 转换为 HWC 格式\n        img = Image.fromarray(img_data)\n        \n        # 使用 transform 进行数据增强或标准化\n        if self.transform:\n            img = self.transform(img)\n\n        label = self.labels[idx]\n        return img, label\n\n\n # 设置路径，指向 test_batch 文件\ntest_batch_path = '/kaggle/input/2024-zju/test_batch'   \n# 创建测试集 Dataset\ntest_dataset = CIFAR10TestDataset(test_batch_path, transform=transform)\n\n# 创建 DataLoader\ntestloader = DataLoader(test_dataset, batch_size=2, shuffle=False)\n\n\nnet = create_model('vit_base_patch16_224', pretrained=False, num_classes=10)  # Ensure the same structure is used\nnet.to(device)\n\n# Load saved model weights\nnet.load_state_dict(torch.load('/kaggle/input/vit/pytorch/default/1/cifar10_vit.pth'))\nnet.eval()  # Set model to evaluation mode\n\n\ncorrect = 0\ntotal = 0\n\nwith torch.no_grad():\n    for data in tqdm(testloader):\n        images, labels = data\n        images, labels = images.to(device), labels.to(device)\n        outputs = net(images)\n        _, predicted = torch.max(outputs.data, 1)\n        total += labels.size(0)\n        correct += (predicted == labels).sum().item()\n\nprint(f'Accuracy: {100 * correct / total:.2f}%')","metadata":{"execution":{"iopub.status.busy":"2024-09-22T03:41:34.196885Z","iopub.execute_input":"2024-09-22T03:41:34.197815Z","iopub.status.idle":"2024-09-22T03:43:08.366750Z","shell.execute_reply.started":"2024-09-22T03:41:34.197771Z","shell.execute_reply":"2024-09-22T03:43:08.365680Z"},"trusted":true},"execution_count":null,"outputs":[]}]}