Accelerating AI Development Pipelines for Industrial Inspection with the NVIDIA TAO Toolkit

This post explores how NVIDIA TAO Toolkit can quickly and accurately train AI models, showing how AI and transfer learning can transform how image and video analysis and industrial processes are deployed.

There is an increasing demand for manufacturers to achieve high-quality control standards in their production processes. Traditionally, manufacturers have relied on manual inspection to guarantee product quality. However, manual inspection is expensive, often only covers a small sample of production, and ultimately results in production bottlenecks, lowered productivity, and reduced efficiency.

By automating defect inspection with AI and computer vision, manufacturers can revolutionize their quality control processes. However, one major obstacle stands between manufacturers and full automation. Building an AI system and production-ready application is hard and typically requires a skilled AI team to train and fine-tune the model. The average manufacturer does not employ this expertise and resorts to using manual inspection.

The goal of this project was to show how the NVIDIA TAO Toolkit and a pretrained model can be used to quickly build more accurate quality control in the manufacturing process. This project was done without an army of AI specialists or data scientists. To see how effective TAO Toolkit is in training an AI system for commercial quality-control purposes, a publicly available dataset on the steel welding process was used to retrain a pretrained ResNet-18 model, from the NGC catalog, a GPU-optimized hub of AI and HPC software, using TAO Toolkit. We compared the effort and resulting model’s accuracy to a model built from scratch on the dataset in a previously published work by a team of AI researchers.

TAO Toolkit is user-friendly and fast, and can be easily used by engineers who do not have AI expertise. We observed that TAO Toolkit was faster to set up and produced more accurate results, posting a macro average F1 score of 97% compared to 78% from previously published “built from scratch” work on the dataset.  

This post explores how TAO Toolkit can quickly and accurately train AI models, showing how AI and transfer learning can transform how image and video analysis and industrial processes are deployed.

Workflow with NVIDIA TAO Toolkit

NVIDIA TAO Toolkit, a core component of the NVIDIA Train, Adapt, and Optimize (TAO) platform, follows the zero-coding paradigm to fast-track AI development. The TAO Toolkit comes with a set of ready-to-use Jupyter notebooks, Python scripts, and configuration specifications with default parameter values that enable you to start training and fine-tuning your datasets quickly and easily.

To get started with the TAO Toolkit, we followed these Quick Start Guide instructions.

  • We downloaded a Docker container and TAO Toolkit Jupyter notebook.
  • We mapped our dataset onto the Docker container.
  • We started our first training, having adjusted the default training parameters, such as the network structure, network size, optimizer, and so on, until we were satisfied with the results.

For more information about the configuration options for training, see Model Config in the NVIDIA TAO Toolkit User Guide.

Dataset

The dataset used in this project was created by researchers at the University of Birmingham for their paper, Automated defect classification of SS304 TIG welding process using visible spectrum camera and machine learning.

The dataset consists of over 45K grayscale welding images, which can be obtained through Kaggle. The dataset describes one class of proper execution: good_weld. It has five classes of defects that can occur during a tungsten inert gas (TIG) welding process: burn_through, contamination, lack_of_fusion, lack_of_shielding_gas, and high_travel_speed.

Figure 1. Sample welding images from the training dataset. Source: Bacioiu et. al, 2019.

 burn_throughcontaminationgood_weldhigh_travel_speedlack_of_fusionlack_of_shieldingTotaltrain977161315752630503619624204valid646339670034615611029694test7319607628249149010211160Table 1. Image distribution across the train, validation, and test data set

Like many industrial datasets, this dataset is quite imbalanced as it can be difficult to collect data on defects that occur with a low likelihood. Table 1 shows the class distribution for the train, validation, and test datasets.

Figure 2 visualizes the imbalance in the test dataset. The test dataset contains 75x more images of good_weld than it does of lack_of_shielding.

Figure 2. Class distribution of the test dataset on TIG steel welding.

Using NVIDIA TAO Toolkit

The approach taken focuses on minimizing both development time and tuning time while ensuring that the accuracy is applicable for a production environment. TAO Toolkit was used in combination with the standard configuration files shipped with the example notebooks. The setup, training, and tuning was done in under 8 hours.

We conducted parameter sweeps regarding network depth and the number of training epochs. We observed that changing the learning rate from its default did not improve the results, so we did not investigate this further and left it at the default. The best results were obtained with a pretrained ResNet-18 model obtained from the NGC catalog after 30 epochs of training with a learning rate of 0.006.

Check out the step-by-step approach in krygol/304SteelWeldingClassification GitHub repo.

MetricPrecisionRecallF1-scoreburn_through0.980.860.92contamination0.990.930.96good_weld0.991.000.99high_travel_speed1.001.001.00lack_of_fusion0.951.000.98lack_of_shielding_gas1.000.991.00micro average0.980.980.98macro average0.980.960.97Table 2. Results attained with a pretrained ResNet-18 after training for 30 epochs with a learning rate of 0.006.

The obtained results were comparably good across all classes. Some lack_of_fusion gas images got misclassified as burn_through and contamination images. This effect was also observed when training the deeper ResNet50, which was even more prone to misclassifying lack_of_fusion as another defective class.

Comparison to the original approach

The researchers at the University of Birmingham chose a different AI workflow. They manually prepared the dataset to lessen the imbalance by undersampling it. They also rescaled their images to different sizes and chose custom network structures.

They used a fully connected neural network (Fully-con6), neural network with two hidden layers. They also implemented a convolutional neural network (Conv6) with three convolutional layers each followed by a max pooling layer and a fully connected layer as the final hidden layer. They did not use skip connections as ResNet does.

The results obtained with TAO Toolkit are even more impressive when compared to results of the custom implementation by the researchers at the University of Birmingham.

 Conv6Fully-con6TAO-ResNet-18MetricF1-scoreF1-scoreF1-scoregood_weld0.990.970.99burn_through0.980.170.92contamination0.900.790.96lack_of_fusion0.940.940.97lack_of_shielding0.00.381.00high_travel_speed0.870.121.00Macro average0.780.560.97Table 3. Comparison of the custom networks and the TAO Toolkit ResNet-18.

Conv6 performed better on average with a macro average F1 of 0.78 but failed completely at recognizing lack_of_shielding gas defects. Fully-con6 performed worse on average, with a macro average F1 of 0.56. Fully-con6 could classify some of the lack_of_shielding gas images but had problems with burn_through and high_travel_speed images. Both Fully-con6 and Conv6 had distinct weaknesses that would prohibit them from being qualified as production-ready. 

The best F1-scores for every class are marked in green in the table. As you can see, the ResNet-18 model trained by TAO Toolkit provided better results with a macro average of 0.97.

Conclusion

We had a great experience with TAO Toolkit, which in general was user-friendly and effective. It was fast to set up, easy to use, and produced results acceptable for production within a short computational time. Based on our experience, we believe that TAO Toolkit provides a great advantage for engineers who are not AI experts but wish to use AI in their production environments. Using TAO Toolkit in a manufacturing environment to automate quality control does not come at a performance cost and the application can often be used with the default settings and a few minor tweaks to outperform custom architectures.

This exploration of using TAO Toolkit to quickly and accurately train AI models shows that there is great potential for AI in industrial processes.

Video. Project walkthrough (Source: Boston Virtual HPC Roadshow)

Acknowledgments

Thanks and gratitude to the researchers—Daniel Bacioiu, Geoff Melton, Mayorkinos Papaelias, and Rob Shaw—for their paper, Automated defect classification of SS304 TIG welding process using visible spectrum camera and machine learning, and publishing their data on Kaggle. Thanks to Ekaterina Sirazitdinova for her help during the development process of this project.

Source:: NVIDIA