examples/demo_jax_distributed.py/make_backbone
def make_backbone():
    return keras.Sequential(
        [
            keras.layers.Rescaling(
                1.0 / 255.0
            ),  # input images are in the range [0, 255]
            keras.layers.Conv2D(
                filters=12, kernel_size=3, padding="same", use_bias=False
            ),
            keras.layers.BatchNormalization(scale=False, center=True),
            keras.layers.Activation("relu"),
            keras.layers.Conv2D(
                filters=24,
                kernel_size=6,
                padding="same",
                use_bias=False,
                strides=2,
            ),
            keras.layers.BatchNormalization(scale=False, center=True),
            keras.layers.Activation("relu"),
            keras.layers.Conv2D(
                filters=32,
                kernel_size=6,
                padding="same",
                use_bias=False,
                strides=2,
                name="large_k",
            ),
            keras.layers.BatchNormalization(scale=False, center=True),
            keras.layers.Activation("relu"),
        ],
        name="backbone",
    )