Training Neural Network in TensorFlow

Training Neural Networks in TensorFlow is a very necessary step because, without training, a model will not learn anything; it will be just a model.

So, in this TensorFlow tutorial, I will define what model training means in TensorFlow, how to train the model by calling the fit() method on the TensorFlow model, and also what the arguments passed to fit().

Finally, I will demonstrate through an example where you will learn to create a tensorflow model from scratch, how to train it and make predictions on the new dataset. This tutorial is beginner-level, but a more complex tensorflow model works on the same concepts.

Training Neural Network in TensorFlow

What does training neural network in tensorflow mean?

First, you have a dataset, which can be related to anything like classifying the data point, image classification, audio, video, etc.

These data are normalized and prepared in the form to feed to the neural network model. The neural networks model takes these data and learns from the data. It finds the patterns in the data.

To learn the patterns, the neural network has some parameters, such as (weights and biases) that are adjusted based on the dataset provided to it.

The main goal of the neural network is to find the best parameter values so that it can used on the new or unseen dataset to make correct predictions. If you are unaware of the parameter’s weight and bias, I suggest visiting the article How to Build Perceptron in Python.

Training involves an iterative process where the model learns from the example (in the training dataset) to improve its ability to make correct predictions on similar or new data (the testing dataset).

Internally, the model used different learning techniques to get the best parameter values, and you have learned in the article How to Compile Neural Network in Tensorflow how to specify these learning techniques.

So, after specifying these techniques, we need to train the neural network model in tensorflow. Here, I will explain how you can train your neural network model.

To train the model, you call the fit() method on the model with the dataset and parameters on which you want to train the model.

The syntax is given below., y_traing_data, epochs=10, batch_size=32, validation_split=0.2)

Let’s understand how this fit() function works.

First, you compile a neural network model, and to train this model on the dataset, you need to call the fit() method on this model.

READ:  How to Create a Matrix in Python [5 Ways]

The fit() method accepts some arguments, such as:

x_train_data: This is the input data called features of the training examples; it is a numpy array or a TensorFlow tensor. It contains a row representing a data point; each column represents a feature.

y_train_data: It is the corresponding set of labels for the training data (x_train_data). It should also be in numPy and tensorflow tensor.

So, to the model, you provide a dataset in x_train and y_train because x_train contains features and y_train contains labels. For example, your model classifies whether the given fruit is a mango.

So here, your dataset will be as shown below.

dataset = {
	'Weight': [150, 200, 120, 180, 100],
	'Color (0 = Green, 1 = Yello)': [0, 1, 0, 1, 0],
	'Label (1 = Mango, 0 = Not Mango)': [0, 1, 0, 1, 0]

In the above dataset, the x_train_data will have features such as column Weight and Color, and the y_train_data is the label for each feature, such as the Label column.

As you can see, the dataset contains 5 data points. For example, look at the data point where weight is 180, colour is 1, and label is 1; when the model sees this data point, it understands it is mango.

During training, the model learns from these features and the corresponding labels for those feature values.

epochs: It means a complete pass through the entire dataset. For example, the training dataset contains 5 data points; if you specify the epocs=2, then during the entire training process, the model iterates over the entire training dataset 2 times.

batch_size: For each epoch, the training dataset is divided into batches; if you specify the batch_size as 10, then based on the dataset, 10 batches are created, and each batch contains the subset of the dataset.

For example, if you have a training data set containing 50 data points and specify the batch_size equal to 5, it creates five batches, each containing 10 data points (a subset of the data).

In each epoch, it goes through each batch one by one and updates its parameter values.

validation_split: It is a small part of the training dataset you provide to the model. The purpose of the small dataset is to monitor the model’s performance while training it on this unseen small data.

So if you specify the validation_split equal to 0.3, 30% of the data from the dataset will be used for validation purposes to see how well the model generalizes to data it hasn’t seen.

READ:  Python Tkinter add function with examples

The epochs, batch_size, and validation_split are the parameters adjusted to get optimal model training; these are part of the hyperparameter tuning.

After knowing the above, you are ready to train your neural network model.

Creating and Training Neural Network in Tensorflow

For example, you need to create a model to classify if the fruit is bananas oranges.

So you have collected the dataset about both fruits, and this dataset contains two columns: weight, colour, and label.

WeightColour(1 for mix of red and yellow, 0 for yellowLabel (1 for orange, 0 for banana

Here, if the fruit has weight and the colour is 1, which means a mix of red and yellow, it is orange and labelled 1. If the fruit has the colour 0, which means yellow, it is a banana and labelled 0.

Now, we need to convert the above dataset into a feedable format to model.

So, create numerical data using the numpy array because the model works on the numeric data. As shown below, convert the two fruit data into a numpy array.

First, import the necessary library that we will use.

import tensorflow as tf
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Input, Dense
import numpy as np

Create a feature for the model. We have two features, weight and colour, in our dataset, so convert those feature values into a numpy array, as shown below.

x_train_data = np.array([
    [150, 1],  
    [200, 0],  
    [100, 1],  
    [250, 1],  
    [130, 0],  
    [140, 0],  

Also, we have the corresponding label for each data point in the x_train_data, so convert those into numpy, as shown below.

y_train_data = np.array([1, 0, 1, 1, 0, 0])

Next, define the sequential model using the Sequential() function as shown below.

model = tf.keras.models.Sequential()

create input, hidden and output layers using the Dense() function, as shown below.

input_layer = tf.keras.layers.Dense(units=2, activation='relu', input_shape=(2,))
hidden_layer = tf.keras.layers.Dense(units=10, activation= 'relu')
output_layer = tf.keras.layers.Dense(units=1, activation='relu')

Add the above layers to the model using the add() function below.


You have created a model. Specify and configure the model for training or compile the model with the optimizer, loss, and metrics function using the compile() method, as shown below.

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

After compiling the model, let’s train the model on the x_train_data and y_train_data, with epochs = 100 and validation_split = 0.2., y_train_data, epochs=100, validation_split=0.2)
Training Neural Network in TensorFlow

From the above output, the model is trained on the dataset for 100 epochs. After training, you can use the model to make predictions.

READ:  Python Django length filter

First, evaluate the model performance using the code below.

loss, accuracy = model.evaluate(x_train_data, y_train_data)
print(f"Loss: {loss}, Accuracy: {accuracy}")

In the above code, calling the evaluate function on the model with x_train_data and y_train_data, after evaluating the model, it returns two value loss and accuracy, indicating the model performance.

Now, use the model to predict new fruit data, as shown below.

new_fruit = np.array([[150, 1]])
prediction = model.predict(new_fruit)
print("Predicted class:", "Bananas" if prediction[0][0] > 0 else "Orange")
Training Neural Network in TensorFlow and Making Prediction

From the above picture, when the model is evaluated, its accuracy is 0.5, which means 50%, and loss is 7.712, which is quite high. That indicates our model needs more training, but we have very little data here, so the model cannot learn properly, so the loss is high.

In the real world, a large amount of data with different variations is fed to a neural network model, which learns from that dataset. Here, the goal is to minimise the loss and increase the accuracy so the model can make correct predictions on the new similar dataset.

Also, we have tested the model with a new_fruit value equal to (150,1); you can see it predicted it is orange; although the model is not well trained, the part of this tutorial is to show you how you can train the neural network model.

Also, remember you can fine-tune or change the parameter value of epochs, etc, on different values to increase the accuracy and decrease the loss.

When training a model in the real world, always provide enough data to model for training, and ensure that you minimise the loss error and maximize the model’s accuracy. This is the parameter through which we can measure the model’s performance.

Also, you can compile the model with different optimizers, losses, and metrics.


You learned how to train the model in TensorFlow using the fit() method on the model and also learned how to pass the dataset to the model as x_training and y_train.

Then, you learned how to divide the dataset into subsets using the batch_size parameter.

You may like to read: