CNN-Scene-Categorization/cnn_categorization_improved.py at main · COMPRyanLI/CNN-Scene-Categorization · 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
from torch import nn
def cnn_categorization_improved(netspec_opts):
    """
 Constructs a network for the improved categorization model.

 Arguments
 --------
 netspec_opts: (dictionary), the improved network's architecture.

 Returns
 -------
 A categorization model which can be trained by PyTorch
   """
    net = nn.Sequential()

    # add layers as specified in netspec_opts to the network
    cov_num = 1
    BN_num = 1
    relu_num = 1
    pooling_num = 1
    drop_num = 1
    in_channels = 3 # equal to 1 if grayscale
    for i in range(len(netspec_opts['kernel_size'])):
            layer_type = netspec_opts['layer_type'][i]

            if layer_type == 'conv':
                # Calculate padding based on kernel size to maintain spatial dimensions
                if isinstance(netspec_opts['kernel_size'][i], int):
                    pad = (netspec_opts['kernel_size'][i] - 1) // 2
                else:
                    pad = ((netspec_opts['kernel_size'][i][0] - 1)//2, (netspec_opts['kernel_size'][i][1] - 1)//2)

                # Add convolutional layer
                net.add_module(f'conv{cov_num}',
                            nn.Conv2d(in_channels=in_channels,
                                        out_channels=netspec_opts['num_filters'][i],
                                        kernel_size=netspec_opts['kernel_size'][i],
                                        stride=netspec_opts['stride'][i],
                                        padding=pad))

                # Update in_channels for next convolutional layer
                in_channels = netspec_opts['num_filters'][i]
                cov_num += 1

            elif layer_type == 'bn':
                # Add batch normalization layer
                net.add_module(f'bn{BN_num}',
                            nn.BatchNorm2d(num_features=netspec_opts['num_filters'][i]))
                BN_num += 1

            elif layer_type == 'relu':
                # Add ReLU activation layer
                net.add_module(f'relu{relu_num}', nn.ReLU())
                relu_num += 1

            elif layer_type == 'pool':
                # Add average pooling layer
                net.add_module(f'pool{pooling_num}',
                            nn.AvgPool2d(kernel_size=netspec_opts['kernel_size'][i],
                                            stride=netspec_opts['stride'][i],
                                            padding=0))
                pooling_num += 1
            elif layer_type == 'drop':
                # Add dropout layer
                net.add_module(f'drop{drop_num}',
                            nn.Dropout(p=0.5))
                drop_num += 1

    net.add_module('gap',  nn.AdaptiveAvgPool2d((1,1)))  # → (N,16,1,1)
    net.add_module('flatten', nn.Flatten())
    net.add_module('fc', nn.Linear(netspec_opts['num_filters'][-1], 16)) # fully connected layer for final classification

    return net