MLAB - Source Code

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
from ast import Str
import tensorflow as tf

def residual_block(x, filters, stride=1):
    """ResNet residual block / building block. Used for ResNet 18 and 34.

    Args:
        x (tf.Keras.Layer): Previous Keras layer
        filters (int): The dimensionality of the output space (i.e. the number of output filters in the convolution).
        stride (int, optional): An integer, specifying the strides of the convolution along the height and width. Defaults to 1.

    Returns:
        tf.Keras.Layer: The model with an residual block at the top. 
    """

    skip = tf.keras.layers.Conv2D(filters, kernel_size=(3,3), strides=(stride,stride), padding='same')(x)
    skip = tf.keras.layers.BatchNormalization()(skip)
    skip = tf.keras.layers.Activation(tf.nn.relu)(skip)

    skip = tf.keras.layers.Conv2D(filters, kernel_size=(3,3), padding='same')(skip)
    skip = tf.keras.layers.BatchNormalization()(skip)

    # Option B - See paper Deep Residual Learning for Image Recognition
    if stride != 1:
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv2D(filters, kernel_size=(1,1), strides=(stride,stride), padding='valid')(x)
    
    out = tf.keras.layers.Add()([x,skip])
    out = tf.keras.layers.BatchNormalization()(out)
    out = tf.keras.layers.Activation(tf.nn.relu)(out)

    return out

def bottleneck_block(x, filters, stride=1, option_b=False):
    """ResNet bottleneck building block. Used for ResNet 50, 101 and 152. 

    Args:
        x (tf.Keras.Layer): Previous Keras layer
        filters (int): The dimensionality of the output space (i.e. the number of output filters in the convolution).
        stride (int, optional): An integer, specifying the strides of the convolution along the height and width. Defaults to 1.

    Returns:
        tf.Keras.Layer: The model with an bottleneck building block at the top. 
    """

    skip = tf.keras.layers.Conv2D(filters, kernel_size=(1,1), strides=(stride,stride), padding='same')(x)
    skip = tf.keras.layers.BatchNormalization()(skip)
    skip = tf.keras.layers.Activation(tf.nn.relu)(skip)

    skip = tf.keras.layers.Conv2D(filters, kernel_size=(3,3), padding='same')(skip)
    skip = tf.keras.layers.BatchNormalization()(skip)
    skip = tf.keras.layers.Activation(tf.nn.relu)(skip)

    skip = tf.keras.layers.Conv2D(filters*4, kernel_size=(1,1), padding='same')(skip)
    skip = tf.keras.layers.BatchNormalization()(skip)

    # Option B - See paper Deep Residual Learning for Image Recognition
    if option_b:
        x = tf.keras.layers.Conv2D(filters*4, kernel_size=(1,1), strides=(stride,stride), padding='valid')(x)
        x = tf.keras.layers.BatchNormalization(axis=3)(x)

    out = tf.keras.layers.Add()([x,skip])
    out = tf.keras.layers.Activation(tf.nn.relu)(out)
    
    return out
    

def resnet50_vitis(input_tensor=None, include_top=True, weight_path=None, return_tensor=False, classes=1000, classifier_activation="softmax"):
    """Creates and returns the ResNet50 CNN architecture.

    Args:
        input_tensor: optional keras layer, like an input tensor.
        include_top: whether to include the top layers or top.
        weight_path: If not none, these weights will be loaded.
        return_tensor: Whether to return the network as tensor or as `tf.keras.model` (if true, weights will not be loaded).
        classes: By default the number of classes are 1000 (ImageNet). Only important `include_top=True`.
        classifier_activation: By default softmax (ImageNet). Only important if `include_top=True`.

    Returns:
        The CNN architecture as `tf.keras.model` if `return_tensor=False`, otherwise as `tf.keras.layers`.
    """

    if input_tensor is None:
        input_tensor = tf.keras.layers.Input(shape=(224,224,3))

    x = tf.keras.layers.ZeroPadding2D()(input_tensor)
    x = tf.keras.layers.Conv2D(64, kernel_size=(7, 7), strides=(2, 2))(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation(tf.nn.relu)(x)
    x = tf.keras.layers.ZeroPadding2D()(x)
    x = tf.keras.layers.MaxPool2D((3, 3), strides=(2, 2), padding='same')(x)

    # conv2_x
    x = bottleneck_block(x, 64, option_b=True)
    x = bottleneck_block(x, 64)
    x = bottleneck_block(x, 64)

    # conv3_x
    x = bottleneck_block(x, 128, 2, option_b=True)
    x = bottleneck_block(x, 128)
    x = bottleneck_block(x, 128)
    x = bottleneck_block(x, 128)

    # conv4_x
    x = bottleneck_block(x, 256, 2, option_b=True)
    x = bottleneck_block(x, 256)
    x = bottleneck_block(x, 256)
    x = bottleneck_block(x, 256)
    x = bottleneck_block(x, 256)
    x = bottleneck_block(x, 256)

    # conv5_x
    x = bottleneck_block(x, 512, 2, option_b=True)
    x = bottleneck_block(x, 512)
    x = bottleneck_block(x, 512)
    
    if include_top is True:
        x = tf.keras.layers.GlobalAveragePooling2D()(x)
        x = tf.keras.layers.Flatten()(x)
        x = tf.keras.layers.Dense(classes, activation=classifier_activation, name="predictions")(x)

    if return_tensor:
        return x
    
    model = tf.keras.Model(input_tensor, x, name="resnet50")
    if weight_path is not None:
        model.load_weights(weight_path)

    return model