Flower_Classification_CNN/model.py at master · omerpanay/Flower_Classification_CNN · GitHub

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
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models

class CustomCNN(nn.Module):
    def __init__(self, num_classes=5):
        super(CustomCNN, self).__init__()
        # Feature extraction layers
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(256)
        self.conv5 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.bn5 = nn.BatchNorm2d(512)

        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.5)

        # Calculate the size of flattened features
        self._to_linear = None
        self._get_conv_output_size()

        # Classifier layers
        self.fc1 = nn.Linear(self._to_linear, 512)
        self.fc2 = nn.Linear(512, num_classes)

        # Feature maps for visualization
        self.feature_maps = {}

    def _get_conv_output_size(self):
        # Create a dummy input to calculate the size
        x = torch.randn(1, 3, 224, 224)
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = self.pool(F.relu(self.bn4(self.conv4(x))))
        x = self.pool(F.relu(self.bn5(self.conv5(x))))
        self._to_linear = x.shape[1] * x.shape[2] * x.shape[3]

    def _hook_features(self, name):
        def hook(module, input, output):
            self.feature_maps[name] = output.detach()
        return hook

    def forward(self, x):
        # Register hooks for visualization
        self.conv1.register_forward_hook(self._hook_features('conv1'))
        self.conv3.register_forward_hook(self._hook_features('conv3'))
        self.conv5.register_forward_hook(self._hook_features('conv5'))

        # Forward pass
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = self.pool(F.relu(self.bn4(self.conv4(x))))
        x = self.pool(F.relu(self.bn5(self.conv5(x))))

        x = x.view(-1, self._to_linear)
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.fc2(x)
        return x

class VGG16FeatureExtractor(nn.Module):
    def __init__(self, num_classes=5):
        super(VGG16FeatureExtractor, self).__init__()
        # Load pretrained VGG16
        vgg16 = models.vgg16(pretrained=True)

        # Freeze all layers
        for param in vgg16.parameters():
            param.requires_grad = False

        # Remove the classifier
        self.features = vgg16.features

        # Add custom classifier
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

class VGG16FineTuned(nn.Module):
    def __init__(self, num_classes=5):
        super(VGG16FineTuned, self).__init__()
        # Load pretrained VGG16
        vgg16 = models.vgg16(pretrained=True)

        # Freeze first two blocks (first 10 layers)
        for i, param in enumerate(vgg16.features.parameters()):
            if i < 10:  # Freeze first two blocks
                param.requires_grad = False

        self.features = vgg16.features
        self.feature_maps = {}

        # Replace classifier
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def _hook_features(self, name):
        def hook(module, input, output):
            self.feature_maps[name] = output.detach()
        return hook

    def forward(self, x):
        # Register hooks for visualization
        self.features[0].register_forward_hook(self._hook_features('conv1'))
        self.features[10].register_forward_hook(self._hook_features('conv3'))
        self.features[20].register_forward_hook(self._hook_features('conv5'))

        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x