Created using Colaboratory

kuanyshbakytuly · Dec 4, 2022 · 3e6dc8e · 3e6dc8e
1 parent 83a8e96
commit 3e6dc8e
Showing 1 changed file with 27 additions and 12 deletions.
diff --git a/main.ipynb b/main.ipynb
@@ -136,17 +136,32 @@
       },
       "outputs": [],
       "source": [
-        "def get_model():\n",
-        "    model = models.vgg16(pretrained=True)\n",
-        "    for param in model.parameters():\n",
+        "def get_model(model_name='VGG16'):\n",
+        "    if model_name == 'VGG16':\n",
+        "      model = models.vgg16(pretrained=True)\n",
+        "      for param in model.parameters():\n",
         "        param.requires_grad = False\n",
-        "    model.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))\n",
-        "    model.classifier = nn.Sequential(nn.Flatten(),\n",
-        "    nn.Linear(512, 128),\n",
-        "    nn.ReLU(),\n",
-        "    nn.Dropout(0.2),\n",
-        "    nn.Linear(128, 1),\n",
-        "    nn.Sigmoid())\n",
+        "      model.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))\n",
+        "      model.classifier = nn.Sequential(nn.Flatten(),\n",
+        "      nn.Linear(512, 128),\n",
+        "      nn.ReLU(),\n",
+        "      nn.Dropout(0.2),\n",
+        "      nn.Linear(128, 1),\n",
+        "      nn.Sigmoid())\n",
+        "    \n",
+        "    if model_name == 'ResNet38':\n",
+        "      model = models.resnet38(pretrained=True)\n",
+        "      for param in model.parameters():\n",
+        "          param.requires_grad = False\n",
+        "      model.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))\n",
+        "      model.fc = nn.Sequential(nn.Flatten(),\n",
+        "      nn.Linear(512, 128),\n",
+        "      nn.ReLU(),\n",
+        "      nn.Dropout(0.2),\n",
+        "      nn.Linear(128, 1),\n",
+        "      nn.Sigmoid())\n",
+        "\n",
+        "\n",
         "\n",
         "    #loss_function - Binary Cross Entropy\n",
         "    #optimizer - Adam\n",
@@ -212,7 +227,7 @@
       "outputs": [],
       "source": [
         "trn_dl, val_dl = get_data()\n",
-        "model, loss_fn, optimizer = get_model()"
+        "model, loss_fn, optimizer = get_model(model='VGG16') # VGG16 for accuracy or ResNet38 for quickly training"
       ]
     },
     {
@@ -401,7 +416,7 @@
     "colab": {
       "provenance": [],
       "mount_file_id": "1vhxvWcfW-qM5DiYNadxexVNyUnCBRCxR",
-      "authorship_tag": "ABX9TyNyuCuyZUaAn8Sa78hxszP2",
+      "authorship_tag": "ABX9TyMj0wyaWT3fjMQUuQC6DRho",
       "include_colab_link": true
     },
     "kernelspec": {