File size: 6,308 Bytes
254c4fb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### train a binary classifier for CIFAKE"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"d:\\Anaconda\\envs\\ai\\Lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
" warnings.warn(\n",
"d:\\Anaconda\\envs\\ai\\Lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.\n",
" warnings.warn(msg)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1/10\n",
"Train Loss: 0.4615 Acc: 0.7786\n",
"Epoch 2/10\n",
"Train Loss: 0.3657 Acc: 0.8369\n",
"Epoch 3/10\n",
"Train Loss: 0.3367 Acc: 0.8522\n",
"Epoch 4/10\n",
"Train Loss: 0.3150 Acc: 0.8626\n",
"Epoch 5/10\n",
"Train Loss: 0.3046 Acc: 0.8680\n",
"Epoch 6/10\n",
"Train Loss: 0.2917 Acc: 0.8745\n",
"Epoch 7/10\n",
"Train Loss: 0.2805 Acc: 0.8809\n",
"Epoch 8/10\n",
"Train Loss: 0.2760 Acc: 0.8824\n",
"Epoch 9/10\n",
"Train Loss: 0.2714 Acc: 0.8843\n",
"Epoch 10/10\n",
"Train Loss: 0.2653 Acc: 0.8868\n",
"Model saved to vgg_model.pth\n",
"Test Accuracy: 0.8859\n",
"Final test accuracy: 0.8859\n"
]
}
],
"source": [
"import os\n",
"import torch\n",
"import torch.nn as nn\n",
"from torch.utils.data import DataLoader, random_split\n",
"from torchvision import datasets, models, transforms\n",
"\n",
"def train_and_save_vgg_model(data_dir, model_path, num_epochs=30, batch_size=32, lr=0.00001):\n",
" data_transform = transforms.Compose([\n",
" transforms.RandomResizedCrop(32),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.RandomRotation(15),\n",
" transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
" ])\n",
"\n",
" full_dataset = datasets.ImageFolder(data_dir, transform=data_transform)\n",
"\n",
" train_size = int(0.8 * len(full_dataset))\n",
" test_size = len(full_dataset) - train_size\n",
" train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])\n",
"\n",
" train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n",
" test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n",
"\n",
" device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
"\n",
" model = models.vgg16(pretrained=True)\n",
" for param in model.features.parameters(): \n",
" param.requires_grad = True\n",
" model.classifier[6] = nn.Linear(model.classifier[6].in_features, 2)\n",
" model = model.to(device)\n",
"\n",
" criterion = nn.CrossEntropyLoss()\n",
" optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
" scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
"\n",
" for epoch in range(num_epochs):\n",
" print(f\"Epoch {epoch + 1}/{num_epochs}\")\n",
" model.train()\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" for inputs, labels in train_loader:\n",
" inputs, labels = inputs.to(device), labels.to(device)\n",
" optimizer.zero_grad()\n",
"\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
" _, preds = torch.max(outputs, 1)\n",
"\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += torch.sum(preds == labels.data)\n",
"\n",
" epoch_loss = running_loss / train_size\n",
" epoch_acc = running_corrects.double() / train_size\n",
"\n",
" print(f\"Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}\")\n",
" scheduler.step()\n",
"\n",
" torch.save(model.state_dict(), model_path)\n",
" print(f\"Model saved to {model_path}\")\n",
"\n",
" model.eval()\n",
" test_corrects = 0\n",
"\n",
" with torch.no_grad():\n",
" for inputs, labels in test_loader:\n",
" inputs, labels = inputs.to(device), labels.to(device)\n",
"\n",
" outputs = model(inputs)\n",
" _, preds = torch.max(outputs, 1)\n",
" test_corrects += torch.sum(preds == labels.data)\n",
"\n",
" test_acc = test_corrects.double() / test_size\n",
" print(f\"Test Accuracy: {test_acc:.4f}\")\n",
"\n",
" return test_acc.item()\n",
"\n",
"\n",
"if __name__ == \"__main__\":\n",
" data_dir = \"CIFAKE/train\" \n",
" model_path = \"vgg_model.pth\"\n",
" num_epochs = 10\n",
" batch_size = 32\n",
" lr = 0.00001\n",
"\n",
" test_acc = train_and_save_vgg_model(data_dir, model_path, num_epochs, batch_size, lr)\n",
" print(f\"Final test accuracy: {test_acc:.4f}\")\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "ai",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.2"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|