ConvNets-TensorFlow2/models/VGG.py at master · archsyscall/ConvNets-TensorFlow2 · 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
import tensorflow as tf
from tensorflow.keras import layers, Sequential, Model

cfg = {
    'A': [64,     'M', 128,      'M', 256, 256,           'M', 512, 512,           'M', 512, 512,           'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256,           'M', 512, 512,           'M', 512, 512,           'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256,      'M', 512, 512, 512,      'M', 512, 512, 512,      'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}


class VGG(Model):
    def __init__(self, features, num_classes, input_shape=(32, 32, 3)):
        super(VGG, self).__init__()

        self.features = Sequential([
            layers.Input(input_shape),
            features
        ])

        self.classifier = Sequential([
            layers.Dense(4096, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(4096, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(num_classes, activation='softmax'),
        ])

    def call(self, inputs, training=False):
        x = self.features(inputs, training=training)
        x = self.classifier(x, training=training)
        return x


def make_layers(cfg):
    nets = []

    for l in cfg:
        if l == 'M':
            nets += [layers.MaxPool2D()]
            continue

        nets += [layers.Conv2D(l, (3, 3), padding='same')]
        nets += [layers.BatchNormalization()]
        nets += [layers.ReLU()]
    return Sequential(nets)


def VGG11(num_classes):
    return VGG(make_layers(cfg['A']), num_classes)


def VGG13(num_classes):
    return VGG(make_layers(cfg['B']), num_classes)


def VGG16(num_classes):
    return VGG(make_layers(cfg['D']), num_classes)


def VGG19(num_classes):
    return VGG(make_layers(cfg['E']), num_classes)