A premade estimator limits the full potential to which TensorFlow can be used; for example, we would not be able to have different dropout values after different layers. In this regard, let's go ahead and create a function of our own, as follows:

Let's explore each part of the preceding snippet of code in detail:

The function takes features (independent variables) and labels (dependent variable) as input. mode indicates whether we want to train, predict, or evaluate the given data.

params provides us with the functionality to supply information about parameters; for example, learning rate:

The preceding snippet of code is similar to the way in which we defined model architecture in Keras, where we specified the inputs, the hidden layer activation, and the number of units in the hidden layer:

If our mode is to predict the class, we would not have to train the model, but just pass the predicted class, thus estimator spec in such scenario would just need to calculate the y_pred_cls values, thus the following code:

If the mode is to train or test the model, we would have to calculate the loss and hence the following code:

In the preceding code, the first line is used to define the cross-entropy calculation. The second line takes the average of cross entropy across all the rows.

optimizer specifies the optimizer we are interested in and the learning rate. train_op specifies that we are interested in minimizing loss, and the global_step parameter keeps a count of the step (epoch) that the model is currently in. metrics specifies the metrics that we are interested in calculating, and the final spec that would be calculated would be a combination of all the preceding parameters that we have defined.

Once the model architecture and the estimator spec that needs to be returned are defined, we define the parameters and mode as follows:

From the preceding code, the function learns the parameter that needs to be changed and also the model architecture that needs to be worked on (model_fn):

We run the model by specifying the mode (in this case, train) by a certain number of epochs (2000 in this case).

After running the model, we evaluate the model's accuracy on the test dataset, as follows: