Customizing Training with TensorFlow in Python
I’ve searched the internet for an easy and complete introduction to custom model building in TensorFlow but I didn’t find it. So I decided to share it and the best reason is, the code explains itself.
Sequential API is a very easy way to create sequential fully connected neural networks. I recommend you to read about it if you don’t know, before continuing this article.
There are several ways to customize your training. Let’s start with custom flows.
1. Custom Flows
You can create custom flows with Functional API. This helps you make models where you have two outputs like MaskRCNN where you have bounding box and class predictions.
For this, you have to use Functional API as shown below.
The base model will not give the best results but the final model will output a bounding box and a classification prediction. You can change the base model as you want.
2. Custom Loss Function
In the training of a neural network, the loss function is very important to find good results. If you want to customize your loss function, I’ll assure you it’s the easiest customization you can make.
To show it, we’ll implement RMSE (root mean squared error), even though it already exists in TensorFlow.
You can change the definition of the function as you want. Don’t forget to pass it to the compile function. I recommend you to check this site if you want to create it as a class.
3. Custom Layers
Here we’ve come to the most important parts.
For this article, we’ll define a quadratic layer wherein for every neuron, we’ll use w0 * x² + w1 * x + b instead of w0 * x + b.
To make it, we define a class that inherits from the Layer class and overwrites three functions.
- The __init__ function will be called when the layer we defined is called. Here we will need to get the variables we want to give to the Layer.
- The build function will be called to create the first state of our variables. Don’t forget to define them as trainable.
- The call function will be called when we want to get the output of this layer. It computes the values that we want to pass to the next layer and returns them.
After that, you can just use this layer just like a Dense layer.
I highly recommend you try to change it as default, where you use the linear function w0 * x + b.
4. Custom Model
Let’s create a custom VGG model.
I’ll first create a Layer that has Max Pooling + k * Convolution + ReLU. Then I’ll create the whole model by repeating it and giving how much convolution layer I want.
After I defined the Block, I’ll create the MyVGG model.
Thus we can customize our model by giving different filters and repetitions to Block or changing the last layers of the VGG. We can customize it as we want, but if you define another layer just don’t forget to call it in the call function otherwise it’ll not change.
And that was the last one. In the last couple of weeks, I’ve finished the TensorFlow: Advanced Techniques Specialization and most of the source code of this article is from the assignments in the course. They have great courses, I highly recommend them to everyone. You can check the all codes from my Github.
Ask me if you have any questions. Please give your feedback. Thank you for reading to the end.
