6
6
7
7
Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
8
8
* 2020-04-03 Initial coding
9
-
9
+ * 2022-12-19 Updated with better comments, improved code using PIL, and checked code still functions as intended.
10
10
"""
11
11
12
12
# Imports
17
17
import torchvision
18
18
import os
19
19
import pandas as pd
20
- from skimage import io
20
+ from PIL import Image
21
21
from torch .utils .data import (
22
22
Dataset ,
23
23
DataLoader ,
@@ -35,7 +35,7 @@ def __len__(self):
35
35
36
36
def __getitem__ (self , index ):
37
37
img_path = os .path .join (self .root_dir , self .annotations .iloc [index , 0 ])
38
- image = io . imread (img_path )
38
+ image = Image . open (img_path )
39
39
y_label = torch .tensor (int (self .annotations .iloc [index , 1 ]))
40
40
41
41
if self .transform :
@@ -50,7 +50,7 @@ def __getitem__(self, index):
50
50
# Hyperparameters
51
51
in_channel = 3
52
52
num_classes = 2
53
- learning_rate = 1e-3
53
+ learning_rate = 3e-4
54
54
batch_size = 32
55
55
num_epochs = 10
56
56
@@ -69,12 +69,19 @@ def __getitem__(self, index):
69
69
test_loader = DataLoader (dataset = test_set , batch_size = batch_size , shuffle = True )
70
70
71
71
# Model
72
- model = torchvision .models .googlenet (pretrained = True )
72
+ model = torchvision .models .googlenet (weights = "DEFAULT" )
73
+
74
+ # freeze all layers, change final linear layer with num_classes
75
+ for param in model .parameters ():
76
+ param .requires_grad = False
77
+
78
+ # final layer is not frozen
79
+ model .fc = nn .Linear (in_features = 1024 , out_features = num_classes )
73
80
model .to (device )
74
81
75
82
# Loss and optimizer
76
83
criterion = nn .CrossEntropyLoss ()
77
- optimizer = optim .Adam (model .parameters (), lr = learning_rate )
84
+ optimizer = optim .Adam (model .parameters (), lr = learning_rate , weight_decay = 1e-5 )
78
85
79
86
# Train Network
80
87
for epoch in range (num_epochs ):
0 commit comments