Training MobileViT on NimbleBox.ai using DeepSpeed

Training a deep learning model, such as the MobileVIT model, can be a complex and challenging task, but it is also a powerful and effective way to solve a wide range of problems in computer vision and other fields. In this blog post, we will discuss the key steps involved in training a MobileVIT model, model definition, and model training.

What is MobileViT?

MobileViT networks are similar to transformers, same technology that powers GPT-3, in that they both use self-attention mechanisms to process input data, but they differ in their architecture and the way they combine the outputs of the self-attention and convolutional modules. Transformers typically use a multi-headed self-attention mechanism, followed by a feedforward neural network, to make predictions, whereas MobileViT networks use a single-headed self-attention mechanism, followed by a convolutional module, to make predictions.

What is DeepSpeed?

DeepSpeed is a deep learning optimization library that is designed to enable training of large-scale deep learning models with minimal code changes. It provides a range of optimization techniques that can be used to improve the performance and efficiency of deep learning models, including model parallelism, data parallelism, mixed precision training, and distributed training. DeepSpeed is built on top of PyTorch and provides a modular, scalable, and user-friendly API that allows users to easily incorporate its optimization techniques into their deep learning workflow. DeepSpeed is open-source and available for use on a wide range of deep learning tasks and applications.

Dataset typically involves collecting a large dataset of images and other data that are relevant to the task that the model will be trained to perform.

NimbleBox Jobs

NBX-Jobs is a job scheduling system that is designed to enable efficient and effective execution of parallel computing tasks on both CPUs and GPUs. It is built on top of the Kubernetes and provides a user-friendly API that allows users to easily define and submit computing jobs to be executed on either CPU or GPU resources. It is better than bare metal VMs because it is easier to use, more scalable, and more resilient, which makes it a better option for deploying and managing applications in many cases.

With nbox scheduling a Job is as simple as running one extra line of code, say you have a train.py file that has a main function like:

def main():
  # load model and all that
  feature_extractor = MobileViTFeatureExtractor.from_pretrained("apple/mobilevit-small")
  model = MobileViTForImageClassification.from_pretrained("apple/mobilevit-small")
  ...
  trainer.train()

All you need to do for this is:

nbx jobs upload train:main --trigger --id '' --resource_cpu="600m" --resource_memory="600Mi" --resource_disk_size="10Gi" --resource_gpu="nvidia-tesla-k80" --resource_gpu_count="1"

To wrap up, NimbleBox Jobs is a critical component of the modern MLOps ecosystem, and it is a platform that every ML developer, operator, and administrator should be familiar with and be able to use effectively. By leveraging the capabilities of NimbleBox Jobs, you can unlock the full potential of your cloud-based ML applications and infrastructure, and take your MLOps to the next level.