Visualize AI Model Training Metrics with TensorBoard

Key Takeaways

in this video we'll show how you can view the training process from the tau toolkit in tensorboard a visual tool for machine learning and experimentation we'll use one of the many computer vision notebooks for this example now i have already configured my system with the tao toolkit and so if you need to do that please refer to the getting started page for more info today i'm going to be using a yolo v4 notebook as an example we are going to look at the impact of changing the learning rate between two versions of the same model learning rate is essentially the size of the steps taken by the model to achieve the lowest cost function to set the learning rate navigate to the specs directory click on the training spec file for any one of these experiments for the first experiment we'll start with the default learning rate as shown here in the second experiment we'll set a higher learning rate we can expect the model to reach the cost function a lot quicker and finally to enable tensorboard set the visualization to true in both the training spec files with the training specification set let's head over to the jupyter notebook to begin our training we'll enable the environment variables by running the first cell in the notebook this is an essential step to enable the training progress the next cells here deal with downloading and installing the tau toolkit launcher which we'll skip as we've already have this installed the next cells here download and convert the format of the data set required for training then we'll download the pre-trained model directly from the ngc catalog and finally you can review the training parameters we said earlier in the cells here okay now we are ready to start training i am going to run the first experiment with default learning rate depending on the size of the model it may take a few minutes for the model to initiate and start solving as you can see that the training process has now started we can now view the progress in tensorboard to do this you'll need to start a new terminal session remotely log into the machine and then navigate to the folder where the results are being written enter the syntax shown here this will enable you to view multiple experiments in the same tensorboard window you can also refer to the documentation for the exact syntax the command will generate a specific url you simply need to copy and paste it into a new browser tab ensure that you enter the correct ip address of the machine so now we have tensorboard running in our browser you can view scalar plots in the first tab these include average precision for the various classes in our model the training loss the learning rate the mean average precision and also the validation loss the next tab shows the model predictions on the input images that are recorded per epoch the next two tabs give you access to some advanced data that includes histograms of the model weights in each layer as well all these metrics are recorded and displayed in real time we'll let this training complete for ade epochs and then view the final set of results the model training is not complete let's take a look at the results we can view the summary of the results in the jupyter notebook now let's take the results in detail in tensorboard as well again you'll see the accuracy for each class within our model we can also see the final training loss the resulting learning rate the mean average precision as well as the validation loss all after 80 epochs so we now know where we stand with our first experiment let's go back to the jupiter notebook and run the case with a higher learning rate again this will take a few minutes to initialize and solve now that this experiment has started to solve let's view the results in tensorboard you can now see the impact of changing the learning rate to higher value the model is converging a lot faster across all the various parameters we'll let this model run through the ad epochs and then inspect our final results you can see quite a big difference in the results between the two experiments shown here in the jupiter notebook it looks like the higher learning rate is showing better results it may just be that we have to run the experiment with the lower learning rate for a few more training epochs now let's compare the results in tensorboard again the comparison plots show that our second experiment the higher learning rate shows better results across the various parameters such as the average precision for each class the training loss the learning rate the mean average precision and also the validation loss with this visual capability we can now see and compare the results of all our experiments in a single view you've just seen how easy and powerful it is to visualize and understand the training process in tensorboard now you can change the right hyper parameters of your model in the tau toolkit experiment with various what-ifs and arrive at an optimized model that fits your use case