human
I am working on the file "00_Miscellaneous/model_optimisation/Tutorial - TensorFlow Model Optimisation for Serving - MNIST.ipynb".
The first task is:
1.2 Estimator
1.2.1 Model Function
Can you write Python code for it?
assistant
#Here is the Python code you requested.

def model_fn(features, labels, mode, params):
    
    is_training = True if mode == tf.estimator.ModeKeys.TRAIN else False

    # convolution layers
    def _cnn_layers(conv_inputs):

        for i in range(params.num_conv_layers):

            filters = params.init_filters * (2**i)
            
            conv = tf.keras.layers.Conv2D(kernel_size=3, filters=filters, strides=1, padding='SAME')(conv_inputs)
            pool = tf.keras.layers.MaxPool2D(pool_size=2, strides=2, padding='SAME')(conv)
            batch_norm = tf.keras.layers.BatchNormalization()(pool, training=is_training)  
            conv_inputs = batch_norm

        outputs = conv_inputs
        return outputs
    
    # fully-connected layers
    def _fully_connected_layers(dense_inputs):
        
        for i in range(len(params.hidden_units)):
            dense = tf.keras.layers.Dense(params.hidden_units[i], activation='relu')(dense_inputs)
            dense_dropout = tf.keras.layers.Dropout(params.dropout)(dense, training=is_training)
            dense_inputs = dense_dropout
            
        outputs = dense_inputs
        return outputs

    # model body
    def _inference(features, mode, params):
        
        input_layer = tf.keras.layers.Reshape([28, 28, 1])(features["input_image"])
        convolutions = _cnn_layers(input_layer)
        flatten = tf.keras.layers.Flatten()(convolutions)
        fully_connected = _fully_connected_layers(flatten)
        
        # unused_layer
        unused_layers = tf.keras.layers.Dense(units=100, name='unused', activation=None)(flatten)
        
        logits = tf.keras.layers.Dense(units=NUM_CLASSES, name='logits', activation=None)(fully_connected)
        return logits
    
    # model head
    head = tf.contrib.estimator.multi_class_head(n_classes=NUM_CLASSES)
    
    return head.create_estimator_spec(
            features=features,
            mode=mode,
            logits=_inference(features, mode, params),
            labels=labels,
           ...