build_resnet.py

# depth should be 9n+2 (eg 56 or 110)

# Model definition
num_filters_in = 32
num_res_block = int((depth - 2) / 9)

inputs = Input(shape=input_shape)

# ResNet V2 performs Conv2D on X before spiting into two path
X = residual_block(X=inputs, num_filters=num_filters_in, conv_first=True)

# Building stack of residual units
for stage in range(3):
    for unit_res_block in range(num_res_block):
        activation = 'relu'
        bn = True
        stride = 1
        # First layer and first stage
        if stage == 0:
            num_filters_out = num_filters_in * 4
            if unit_res_block == 0:
                activation = None
                bn = False
            # First layer but not first stage
        else:
            num_filters_out = num_filters_in * 2
            if unit_res_block == 0:
                stride = 2

        # bottleneck residual unit
        y = residual_block(X,
                           num_filters=num_filters_in,
                           kernel_size=1,
                           stride=stride,
                           activation=activation,
                           bn=bn,
                           conv_first=False)
        y = residual_block(y,
                           num_filters=num_filters_in,
                           conv_first=False)
        y = residual_block(y,
                           num_filters=num_filters_out,
                           kernel_size=1,
                           conv_first=False)
        if unit_res_block == 0:
            # linear projection residual shortcut connection to match
            # changed dims
            X = residual_block(X=X,
                               num_filters=num_filters_out,
                               kernel_size=1,
                               stride=stride,
                               activation=None,
                               bn=False)
        X = tf.keras.layers.add([X, y])
    num_filters_in = num_filters_out