Simple Custom Object classification with Pytorch | ONNX inference
Object classification is a fundamental task in computer vision that involves training a model to classify images into predefined categories. In this blog post, we will explore how to build a simple custom object classification model using PyTorch and then deploy it for inference using the ONNX format. This tutorial assumes a basic understanding of PyTorch and neural networks. Code for my blog is here you can clone it i would go through each of the segment.
- Dataset Preparation: Before diving into the model creation, it’s essential to prepare a labeled dataset. Collect images of different object classes you want to classify and organize them into separate folders based on their categories. Ensure you have a sufficient number of images for each class to avoid overfitting.
Prepare the dataset like below tree
- data
- Fruits (dataset name)
- train
- class 1
- class 2
- ...
- class n
- val
- class 1
- class 2
- ...
- class n
- test
- class 1
- class 2
- ...
- class n- I have took fruits dataset from kaggle for example purpose link here
- change the path name in "train_dir" and "val_dir" in main.py
- Initialize data loaders add augmentations if you need.
2. Building the Model: Start by importing the necessary libraries, including PyTorch. Define the architecture of your custom object classification model, typically using convolutional neural networks (CNNs). Design the network layers, including convolutional and pooling layers
import torch.nn as nn
class CustomConvNet(nn.Module):
def __init__(self, num_classes):
super(CustomConvNet, self).__init__()
self.num_classes = num_classes
self.layer1 = self.conv_module(3, 16)
self.layer2 = self.conv_module(16, 32)
self.layer3 = self.conv_module(32, 64)
self.layer4 = self.conv_module(64, 128)
self.layer5 = self.conv_module(128, 256)
self.gap = self.global_avg_pool(256, self.num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = self.layer5(out)
out = self.gap(out)
out = out.view(-1, self.num_classes)
return out
def conv_module(self, in_num, out_num):
return nn.Sequential(
nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_num),
nn.LeakyReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
def global_avg_pool(self, in_num, out_num):
return nn.Sequential(
nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_num),
nn.LeakyReLU(),
nn.AdaptiveAvgPool2d((1, 1)))3. Training the Model and Evaluating: Implement the training loop, including forward and backward propagation, loss calculation, and gradient optimization using techniques like stochastic gradient descent (SGD) or Adam optimizer. Monitor the training process by tracking metrics like loss and accuracy. We have made use of techniques like data augmentation and regularization to improve model generalization.
python main.pyfor epoch in range(num_epochs):
print("Epoch No -", epoch)
model.train()
running_loss = 0.0
running_corrects = 0
for inputs, labels in dataLoaders["train"]:
# Feeding input and labels to device
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
with torch.set_grad_enabled(True):
outputs = model(inputs)
_, preds = torch.max(outputs,1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()* inputs.size(0)
#calculate accuracy
running_corrects += torch.sum(preds == labels.data)
#scheduler step
exp_lr_scheduler.step()
# Calculate average loss and acc for a epoch
epoch_loss = running_loss/len(train_data)
epoch_acc = running_corrects.double()/len(train_data)
print('Loss:{} , Acc{}'.format(epoch_loss, epoch_acc))
# Saving model every five epoch
if (epoch%5 == 0):
save_model(model,epoch)To run prediction run
Make sure if you have changed weights filename, train folder, input image in predict.py
python predict.py4. Exporting to ONNX Format: Once your model is trained and performs well on the validation set, it’s time to export it to the ONNX format for deployment and inference. ONNX (Open Neural Network Exchange) is an open standard format that allows interoperability between different deep learning frameworks. PyTorch provides tools to export models to ONNX.
Make sure if you have changed the weights filename in export.py
python export.py# Now we will save this model.
import torch.onnx
torch.onnx.export(model,
img,
"./savedModels/custommodel.onnx",
export_params=True,
opset_version=10,
verbose=True, # Print verbose output
input_names=['input'], # Names for input tensor
output_names=['output'])5. Inference with ONNX: Load the saved ONNX model and perform inference on new unseen images. Use the ONNX runtime library to load the model, provide input data, and obtain predictions. Measure the inference time and compare it with the PyTorch model’s inference time to assess any performance improvements achieved through ONNX optimization.
Make sure if you have changed weights filename, train folder, input image in predict.py
python onnx_inference.py# Load the ONNX model
onnx_model = onnx.load("./savedModels/custommodel.onnx")
# Create an ONNX runtime session
ort_session = onnxruntime.InferenceSession("./savedModels/custommodel.onnx")
inputs = {"input": trans_img.numpy()}
outputs = ort_session.run(None, inputs)In this tutorial, we explored the process of building a simple custom object classification model using PyTorch. We learned how to train the model, evaluate its performance, and export it to the ONNX format for inference. By leveraging the ONNX runtime, we demonstrated how to perform inference on new images efficiently. With this knowledge, you can now apply custom object classification to various real-world applications and deploy your models seamlessly.
Credits
