Train Image Classification Model using VIT and Smart Sifting.

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

In this notebook we will train a image classification model using Vision Transformer (VIT) and Smart Sifting library. VIT is a transformer encoder model pretrained on large collection of images from ImageNEt at a resolution of 224X224 pixels.

1.Introduction to Smart Sifting

Smart Sifting is a framework to speed up training of PyTorch models. The framework implements a set of algorithms that filter out inconsequential training examples during training, reducing the computational cost and accelerating the training process. It is configuration-driven and extensible, allowing users to add custom logic to transform their training examples into a filterable format. Smart sifting provides a generic utility for any DNN model, and can reduce the training cost by up to 35% in infrastructure cost.

Smart sifting’s task is to sift through your training data during the training process and only feed the more informative samples to the model. During typical training with PyTorch, data is iteratively sent in batches to the training loop and to accelerator devices (e.g. GPUs or Trainium chips) by the PyTorch data loader. Smart sifting is implemented at this data loading stage and is thus independent of any upstream data preprocessing in your training pipeline. Smart sifting uses your live model and a user specified loss function to do an evaluative forward pass of each data sample as it is loaded. Samples which are high loss will materially impact model training and thus are included in training data; meanwhile data samples which are relatively low loss are already well represented by the model and so are set aside and excluded from training. A key input to smart sifting is the proportion of data to exclude: for example, by setting the proportion to 25%, samples in approximately the bottom quartile of loss of each batch will be excluded from training. Once enough high-loss samples have been identified to complete a batch, the data is sent through the full training loop and the model learns and trains normally. Customers don’t need to make any downstream changes to their training loop when smart sifting is enabled.|

2. Install Required Dependencies

3. Prepare Dataset

For this training we will be using Caltech-101 dataset. Caltech-101 consists of pictures of objects belonging to 101 classes. Each class contains roughly 40 to 800 images, totalling around 9k images. Images are of variable sizes, with typical edge lengths of 200-300 pixels.

Lets start by downloading and extracting the dataset.

We will convert the downloaded data into huggingface datasets arrow format. Note: This is done for convenience not a requirement for sifting library. 

After this step we should have a dataset with train and validation splits. Lets print the dataset to confirm.

Upload Dataset to S3 for Training

Lets upload the dataset to S3 , for this we will leverage the dataset API S3 integration to directly save DataSet object to s3.

4. Run training Job using SageMaker Training.

Adding Sifting library to the Image classification code involves following the below steps

1. Define Loss Function - For Image classification we use CrossEntropy loss defined as below

   class ImageLoss(Loss):
   """
   This is an implementation of the Loss interface for the model 
   required for Smart Sifting. 
   """
   def __init__(self):
       self.celoss = torch.nn.CrossEntropyLoss(reduction='none')
   
   def loss(
           self,
           model: torch.nn.Module,
           transformed_batch: SiftingBatch,
           original_batch: Any = None,
   ) -> torch.Tensor:
       device = next(model.parameters()).device
       batch = {k: v.to(device) for k, v, in original_batch.items()}
   
       # compute loss
       outputs = model(**batch)
       return self.celoss(outputs.logits, batch["labels"])
   
   

2. Define Transformation Function to convert input batch to sifting format.

class ImageListBatchTransform(SiftingBatchTransform):
    """
    This is an implementation of the data transforms for the model 
    required for Smart Sifting. Transform to and from ListBatch
    """
    def transform(self, batch: Any):
        inputs = []
        labels = []

        for i in range(len(batch["pixel_values"])):
            inputs.append(batch["pixel_values"][i])

        for i in range(len(batch["labels"])):
            labels.append(batch["labels"][i])

        return ListBatch(inputs, labels)
    
    def reverse_transform(self, list_batch: ListBatch):
        a_batch = {}
        a_batch["pixel_values"] = self.stack_tensors(list_batch.inputs)
        a_batch["labels"] = self.stack_tensors(list_batch.labels)
  
        return a_batch

    def stack_tensors(self,list_of_tensors):
        if list_of_tensors:
            t = torch.stack(list_of_tensors)
        else:
            t = torch.tensor([])
        return t

3. Define sifting config - Define configuration for sifting.

   Beta_value depicts the proportion of samples to keep , higher the value more samples are sifted. loss_history_length - Depicts the window of samples to include when evaluating relative loss.

   sift_config = RelativeProbabilisticSiftConfig(
            beta_value=3,
            loss_history_length=500,
            loss_based_sift_config=LossConfig(
                 sift_config=SiftingBaseConfig(sift_delay=10)
            )
        )

4. **Wrap the Pytorch Data Loader with Sifting Data loader. As a last step we wrap the Pytroch Dataloader with siftingDataLoader passing config, transformation and loss functions.

  train_dataloader = SiftingDataloader(
                sift_config=sift_config,
                orig_dataloader=train_dataloader,
                batch_transforms=ImageListBatchTransform(),
                loss_impl=ImageLoss(),
                model=model
        )  

We define few metrics to be tracked inorder to monitor sifting. This are optional metrics useful to debug and understand sifting performance.

We will launch the training job using G5.2xlarge instance. Smart Sifting library is part of the SageMaker Pytorch Deep Learning containers starting version 2.0.1.

Launch the training job with Data in S3

In this notebook, we looked at how to use smart sifting library to train an Image classification model. Smart sifting helps in reducing training time upto 40% without any reduction in Model performance.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.