initial commit

README.md

@@ -1 +1,84 @@
-# ML-Engineer-Challenge
+# Machine Learning Engineer Coding Challenge
+
+Welcome! you are a super star for making it here. This is your time to shine, an opportunity to show off your skills, understanding and more importantly coding abilities 😉. So relax, grab some coffee / whiskey (depending on time of day) and start developing on this take-home exercise.
+
+
+## Overview
+
+This coding test is divided into three parts, each testing different aspects of your machine learning engineering skills. You will need to use PyTorch, TensorRT, ONNX, and various hyperparameter tuning libraries to complete these tasks. 
+
+Please ensure you have the necessary libraries installed. If you do not have a GPU environment, please let `[email protected]` know, and one will be created for you.
+
+## Part 1: Model Quantisation and Benchmarking
+
+**Objective**: Take a complex computer vision model from Torch Hub, quantise it, and benchmark the speed of inference on the test subset of [tiny-ImageNet dataset](https://www.kaggle.com/datasets/akash2sharma/tiny-imagenet).
+
+To download the dataset, use the following utility script:
+
+```shell
+python ./utils/download_tiny_imagenet.py
+```
+
+### Instructions 📃
+
+1. Select a complex computer vision model from Torch Hub (e.g., Dinov2).
+2. Prepare a small subset of ImageNet images for inference.
+3. Apply dynamic quantisation to the model.
+4. Measure and compare the inference time of the original and quantized models.
+5. Report the inference times and any differences in accuracy.
+
+### Submission 💻
+- Python script with code for loading, quantising, and benchmarking the model.
+- A brief analysis report (ipynb, markdown or PDF) with inference time comparisons and accuracy differences.
+
+---
+
+## Part 2: Automated Hyperparameter Tuning
+
+**Objective**: Conduct automated hyperparameter tuning to identify the optimal hyperparameters for a small CNN trained on the [tiny-ImageNet dataset](https://www.kaggle.com/datasets/akash2sharma/tiny-imagenet) training dataset. (Refer to instructions in Part 1 for downloading the data)
+
+### Instructions 📃
+
+1. Define a small CNN architecture of choice for the CIFAR-100 dataset.
+2. Set up a training loop for the CNN model.
+3. Choose hyperparameters to tune (e.g., learning rate, batch size, number of layers, etc.).
+4. Use a hyperparameter optimization library (e.g., Optuna, Hyperopt, or Scikit-Optimize) to find the best hyperparameters.
+5. Train the model using the optimal hyperparameters and report the final accuracy.
+
+### Submission 💻
+
+- Python script with the model definition, training loop, and hyperparameter tuning setup.
+- A brief report (markdown or PDF) detailing the hyperparameter tuning process and final model accuracy.
+
+---
+
+## Part 3: Model Conversion to TensorRT and ONNX
+
+Objective: Convert a trained model to TensorRT format and serialize it in ONNX for fast inference on Nvidia GPUs.
+
+### Instructions 📃
+
+1. Train or use a pre-trained model (it can be the model from Part 1 or another model).
+2. Export the model to ONNX format.
+3. Convert the ONNX model to TensorRT using TensorRT tools.
+4. Measure the inference time of the TensorRT model on an Nvidia GPU.
+5. Report the inference times and any speedup achieved.
+
+### Submission 💻
+
+- Python script with code for model training/loading, ONNX export, and TensorRT conversion.
+- A brief report (markdown or PDF) with inference time benchmarks and any observed improvements.
+
+## General Submission Guidelines
+
+- Ensure all code is well-documented and follows best practices.
+- Include a requirements.txt file with all dependencies required to run your code.
+- Submit your code and reports in a zip file or through a GitHub repository link.
+
+## Evaluation Criteria
+- Correctness: Does the code achieve the desired outcomes?
+- Efficiency: Are the implementations optimized for performance?
+- Clarity: Is the code well-structured and documented?
+- Reporting: Are the reports clear and do they adequately explain the results?
+
+Good luck, and we look forward to seeing your solutions!

utils/basic_cnn.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+# Example usage
+if __name__ == "__main__":
+    block = BasicBlock(64, 128, stride=2)
+    print(block)

utils/download_tiny_imagenet.py

@@ -0,0 +1,28 @@
+import os
+import requests
+import zipfile
+
+def download_tiny_imagenet(url, dataset_path='tiny-imagenet-200'):
+    # Define the download path
+    download_path = f'{dataset_path}.zip'
+
+    # Download the dataset
+    print(f'Downloading {url}...')
+    response = requests.get(url, stream=True)
+    with open(download_path, 'wb') as file:
+        for chunk in response.iter_content(chunk_size=128):
+            file.write(chunk)
+
+    # Extract the dataset
+    print(f'Extracting {download_path}...')
+    with zipfile.ZipFile(download_path, 'r') as zip_ref:
+        zip_ref.extractall(dataset_path)
+
+    # Clean up the zip file
+    os.remove(download_path)
+    print(f'Dataset downloaded and extracted to {dataset_path}')
+
+# Example usage
+if __name__ == "__main__":
+    url = 'http://cs231n.stanford.edu/tiny-imagenet-200.zip'
+    download_tiny_imagenet(url)

utils/sample_serving_script.py

@@ -0,0 +1,32 @@
+from flask import Flask, request, jsonify
+import torch
+from torchvision import models, transforms
+from PIL import Image
+import io
+
+app = Flask(__name__)
+model = models.resnet50(pretrained=True)
+model.eval()
+
+def transform_image(image_bytes):
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    image = Image.open(io.BytesIO(image_bytes))
+    return transform(image).unsqueeze(0)
+
+@app.route('/predict', methods=['POST'])
+def predict():
+    if request.method == 'POST':
+        file = request.files['file']
+        img_bytes = file.read()
+        tensor = transform_image(img_bytes)
+        outputs = model(tensor)
+        _, predicted = torch.max(outputs, 1)
+        return jsonify({'class_id': predicted.item()})
+
+if __name__ == '__main__':
+    app.run()

utils/tiny_imagenet_dataloader.py

@@ -0,0 +1,65 @@
+import os
+import requests
+import zipfile
+from torchvision import transforms, datasets
+from torch.utils.data import DataLoader
+
+def download_and_extract_tiny_imagenet(url, dataset_path='tiny-imagenet-200'):
+    # Define the download path
+    download_path = f'{dataset_path}.zip'
+
+    # Download the dataset
+    print(f'Downloading {url}...')
+    response = requests.get(url, stream=True)
+    with open(download_path, 'wb') as file:
+        for chunk in response.iter_content(chunk_size=128):
+            file.write(chunk)
+
+    # Extract the dataset
+    print(f'Extracting {download_path}...')
+    with zipfile.ZipFile(download_path, 'r') as zip_ref:
+        zip_ref.extractall(dataset_path)
+
+    # Clean up the zip file
+    os.remove(download_path)
+    print(f'Dataset downloaded and extracted to {dataset_path}')
+
+def get_tiny_imagenet_dataloaders(data_dir, batch_size=32, num_workers=2):
+    # Define the transforms for training and validation
+    transform_train = transforms.Compose([
+        transforms.RandomResizedCrop(64),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.4802, 0.4481, 0.3975], std=[0.2302, 0.2265, 0.2262]),
+    ])
+
+    transform_val = transforms.Compose([
+        transforms.Resize(64),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.4802, 0.4481, 0.3975], std=[0.2302, 0.2265, 0.2262]),
+    ])
+
+    # Load the datasets with ImageFolder
+    train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transform=transform_train)
+    val_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transform=transform_val)
+
+    # Create DataLoaders
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+    return train_loader, val_loader
+
+# Example usage
+if __name__ == "__main__":
+    # Download and extract the dataset
+    url = 'http://cs231n.stanford.edu/tiny-imagenet-200.zip'
+    dataset_path = 'tiny-imagenet-200'
+    if not os.path.exists(dataset_path):
+        download_and_extract_tiny_imagenet(url, dataset_path)
+
+    # Create DataLoaders
+    train_loader, val_loader = get_tiny_imagenet_dataloaders(dataset_path, batch_size=32, num_workers=4)
+
+    # Print dataset sizes
+    print(f'Training set size: {len(train_loader.dataset)}')
+    print(f'Validation set size: {len(val_loader.dataset)}')