Last Updated on 29/04/2026 by Eran Feit
Why FasterViT? The Power of Hybrid CNN-ViT Architectures Moving beyond standard architectures often feels like a trade-off between speed and accuracy. If you are looking to train FasterViT PyTorch custom dataset models, you’ve likely realized that NVIDIA’s hybrid approach is the current SOTA for throughput. In this guide, we solve the challenge of preparing a non-standard dataset—specifically our Star Wars character collection—and feeding it into the FasterViT pipeline. You’ll learn exactly how to configure the model to leverage both convolutional efficiency and hierarchical attention, transforming raw images into a high-performance classifier that outperforms traditional Vision Transformers.
FasterViT image classification using custom dataset represents a modern, efficient approach to training deep learning models that can recognize and categorize images from your own tailored collection of visual data. In a world where off-the-shelf datasets often don’t match specific application needs, applying models like FasterViT to a custom dataset — such as a curated set of Star Wars character images — allows developers to build computer vision systems that reflect real-world use cases and unique classification requirements. This customization is particularly powerful for niche tasks that require specialized image recognition beyond generic object categories, making FasterViT an exciting choice for advanced computer vision practitioners.
At its core, FasterViT combines the strengths of traditional convolutional neural networks (CNNs), which excel at capturing detailed local features, with the global context modeling of vision transformers. This hybrid architecture enables the model to learn both fine-grained texture patterns and broad relational information across an image, leading to robust feature representations. When applied to a custom dataset, better feature extraction and classification accuracy can be achieved compared to using either a pure CNN or a pure transformer model alone.
Training FasterViT on a custom dataset introduces challenges and opportunities. Developers must carefully prepare the data — organizing training, validation, and testing splits, and ensuring proper image preprocessing — so the model can effectively generalize from limited or imbalanced samples. With proper dataset preparation and training strategies, custom FasterViT models can outperform many conventional deep learning classifiers, especially in tasks that require distinguishing fine differences between similar image classes.
Finally, integrating FasterViT into real-world applications demonstrates its practical value. Whether building a character recognizer from a Star Wars image set, designing a wildlife species classifier, or developing a custom industrial defect detector, FasterViT’s ability to leverage custom datasets makes it a highly adaptable tool in the computer vision ecosystem. As deep learning continues to evolve, models like FasterViT that blend efficiency with performance are key for developers who need both speed and accuracy in specialized image classification tasks.
If you are new to the Transformer landscape, I recommend starting with my Vision Transformer Image Classification PyTorch Tutorial to understand the fundamental architecture before moving to the ‘Faster’ variants.
🌟 What Is FasterViT Image Classification and Why Custom Data Matters? FasterViT image classification is a deep learning approach that blends convolutional neural networks with vision transformer elements to effectively classify images. Unlike traditional models that focus purely on either local patterns (CNNs) or global relationships (transformers), FasterViT uses a hybrid architecture that takes advantage of the best of both worlds. This makes it especially suitable for image classification tasks where details matter and large-scale context influences predictions, such as distinguishing between visually similar characters or objects.
When working with custom datasets — like a unique Star Wars image collection — conventional models often fall short due to limited examples or insufficient context. Custom datasets reflect real-world problems these models need to solve, such as identifying specific character traits or unusual visual features unseen in benchmark datasets. Training FasterViT on these tailored images gives the model exposure to exactly the visual domains it will encounter during inference, improving accuracy and robustness.
Moreover, FasterViT’s hierarchical attention mechanism enables efficient processing of visual features at different scales. Early layers might focus on small, local details in images — edges, textures, or character accessories — while later transformer-like layers capture broader patterns or global relationships across the entire image. This layered attention system supports the model’s ability to generalize from a custom dataset to unseen test images, making it well-suited for specialized vision tasks that vary significantly from popular benchmark collections.
In practical application, this means that when you feed FasterViT a custom dataset of labeled images, the model can learn both subtle and global cues needed to make reliable classifications. Whether your goal is to recognize specific individuals, object categories, or nuanced visual differences dictated by your dataset’s uniqueness, FasterViT’s flexible architecture enables innovative solutions. By training this model on your curated images, you customize the learning process — and the resulting classifier becomes tailored, higher performing, and more relevant to your specific needs.
The FasterViT architecture o begin our FasterViT Image Classification with a Custom Dataset in Python , we must first initialize the model with weights tailored for hierarchical feature extraction. FasterViT represents a breakthrough in vision transformers by using a multi-scale architecture. Unlike standard ViT models, implementing FasterViT Image Classification with a Custom Dataset in Python allows us to leverage carrier tokens that pass information across local windows, significantly reducing the computational cost of global self-attention.
Figure 1: The hierarchical attention mechanism used in our FasterViT Image Classification with a Custom Dataset in Python workflow. The FasterViT architecture is built as a hybrid model that combines the strengths of convolutional neural networks and vision transformers in a single unified pipeline. The network begins with traditional convolutional layers, which downsample the input image and extract low-level visual features such as edges, textures, and simple spatial patterns. These convolution stages are efficient and computationally light, making them ideal for early processing where fine-grained image structure matters most. By progressively reducing the spatial resolution while increasing channel depth, the model prepares compact yet informative feature maps for the next stages.
After the initial convolution blocks, FasterViT transitions into deeper stages that incorporate hierarchical attention. This is where the transformer-based components come into play. Hierarchical attention allows the model to capture long-range dependencies in the image — understanding how different regions relate to one another, even when they are far apart spatially. This global reasoning is what makes transformer-based architectures particularly powerful for image understanding, as the model is no longer limited to only local receptive fields like a CNN. FasterViT carefully balances attention computation so it remains efficient while still modeling complex contextual relationships.
In the later stages, the model continues to alternate between downsampling and attention-driven processing, building increasingly abstract feature representations. By the time the data reaches the classification head, the network has learned both detailed local information and high-level contextual structure. This combination enables FasterViT to achieve strong accuracy on image classification tasks while remaining computationally efficient compared to pure transformer models. The architecture is therefore especially useful for real-world applications where both performance and speed matter, making it a compelling evolution in the vision transformer family.
If your hardware constraints are extremely tight, you might also consider comparing these results against a modern CNN approach by following my guide on How to Train ConvNeXt in PyTorch on a Custom Dataset .
FasterViT animal classification process Building a Practical FasterViT Image Classification Tutorial with a Custom Dataset This tutorial walks through the full process of using FasterViT image classification on a custom dataset , showing how to install the required libraries, prepare the dataset, train the model, and finally test it on new images. The goal is to give you a complete, working pipeline so you can adapt the same approach to any dataset you choose — whether that’s animals, products, vehicles, or any type of labeled image collection. Everything is designed to be hands-on and code-driven so you can follow along step-by-step.
The core idea behind the tutorial is to demonstrate how FasterViT — a hybrid model that combines convolutional layers with transformer-based attention — can be trained on real-world images rather than relying only on standard benchmark datasets. You’ll see how the dataset is split into training, validation, and test sets, how transformations are applied, and how the model learns to identify each class. By the end, you’ll understand not only how the code works, but also why each stage of the workflow is important.
Another key part of the tutorial is modifying the model’s final classification head to match the number of classes in your dataset. This ensures FasterViT can correctly predict the labels you define. The training loop also tracks accuracy and loss over time, so you can monitor how well the model is learning and automatically keep the best-performing weights. This makes the process robust and beginner-friendly, while still being powerful enough for advanced experimentation.
Finally, the tutorial shows how to load your trained model and run predictions on new images. You’ll see how to preprocess test images, send them through the model, and display the predicted class — even overlaying the result on the image itself. This allows you to move from raw data → trained model → working classifier , creating a complete solution that you can reuse and scale for future projects.
Link to the video tutorial : https://youtu.be/n-SpVoHrzDQ
You can download the code for the tutorial here :https://eranfeit.lemonsqueezy.com/checkout/buy/a6159108-c66c-4e21-80e0-7a6589f0b8b0 or here : https://ko-fi.com/s/28ca45253c
Link for to the post for Medium users : https://medium.com/vision-transformers-tutorials/fastervit-image-classification-using-custom-dataset-star-wars-dataset-8e6ce470d566
Master Computer Vision
Follow my latest tutorials and AI insights on my
Personal Blog .
Beginner Complete CV Bootcamp
Foundation using PyTorch & TensorFlow.
Get Started → Interactive Deep Learning with PyTorch
Hands-on practice in an interactive environment.
Start Learning → Advanced Modern CV: GPT & OpenCV4
Vision GPT and production-ready models.
Go Advanced →
FasterViT Image Classification Using Custom Dataset FasterViT is a powerful hybrid model that combines the feature extraction strengths of convolutional neural networks with the global contextual understanding of vision transformers.
The goal is practical mastery: from environment setup, dataset preparation, training, and testing, you will walk through a real-world workflow from top to bottom.
Setting Up Your Environment: Core Requirements for FasterViT and PyTorch Before training FasterViT on your custom dataset, your system must be ready.
You’ll also check the CUDA version to ensure GPU acceleration is available, which significantly speeds up training.
# Create a new Conda environment named "fasterVit" with Python 3.11 conda create - n fasterVit python = 3.11 # Activate the newly created "fasterVit" environment conda activate fasterVit # Check the installed CUDA version on your system nvcc -- version # Install PyTorch 2.5.0, Torchvision, Torchaudio, and CUDA 12.4 support from PyTorch and NVIDIA channels conda install pytorch == 2.5 . 0 torchvision == 0.20 . 0 torchaudio == 2.5 . 0 pytorch - cuda = 12.4 - c pytorch - c nvidia # Install the FasterVit library version 0.9.8 from PyPI pip install fastervit == 0.9 . 8 # Install timm (PyTorch Image Models) version 0.9.12 for model utilities and backbones pip install timm == 0.9 . 12 # Install matplotlib for plotting and visualization pip install matplotlib # Install OpenCV Python bindings (note: command is missing the 'install' verb as written) pip opencv - python == 4.10 . 0 . 84 This environment gives you all the tools needed to implement FasterViT training and testing without friction.
Setting Up Your FasterViT PyTorch Custom Dataset: Structure & Labels Before we start training the FasterViT image classification model, we first need a labeled image dataset. In this tutorial, we are working with a collection of character images that are already grouped into folders by class. Each folder represents one category, and every image inside belongs to that specific label. This structure is very important because PyTorch relies on the folder layout to automatically map images to class names during training.
The dataset contains multiple characters, and each character has a number of image samples from different angles, lighting conditions, and backgrounds. This variation helps the model generalize and learn what truly defines each class, rather than memorizing a single image pattern. If your dataset includes clear, centered subjects and consistent labeling, you will normally achieve stronger and more reliable classification results.
Want the exact dataset so your results match mine?
If you want to reproduce the same training flow and compare your results to mine, I can share the dataset structure and what I used in this tutorial. Send me an email and mention the name of the tutorial / dataset , so I know what you’re requesting.
🖥️ Email: feitgemel@gmail.com
D : / Data - Sets - Image - Classification / Star - Wars - Characters
Inside this folder, each sub-folder represents a different class. For example:
Star - Wars - Characters / ├── Class_1 / ├── Class_2 / ├── Class_3 / ├── ... This means you do not need a CSV file or manual labels — the folder names themselves act as the class labels. Later, our script automatically splits these images into Train, Validation, and Test folders while preserving the class-based structure. If you ever decide to swap this dataset with your own images, just keep the same folder-per-class layout and the rest of the code will continue to work smoothly.
Curating the Star Wars Dataset: Structuring Folders for the PyTorch ImageFolder Pipeline The next step is organizing your custom dataset into training, validation, and testing splits.
This code creates appropriate folder structures, randomly shuffles images, and distributes them into the right folders.
When working with Transformers like FasterViT, your preprocessing steps must strictly match the model’s expected patch size (typically 224×224). Beyond just resizing, ensure you are using the specific normalization statistics (mean and standard deviation) from the pre-trained weights. This alignment prevents ‘feature drift’ during training and significantly accelerates the convergence of your custom Star Wars classifier.
Proper data augmentation is the backbone of any FasterViT Image Classification with a Custom Dataset in Python project. Because vision transformers are data-hungry, we use torchvision.transforms to artificially expand our Star Wars dataset. When you are performing FasterViT Image Classification with a Custom Dataset in Python , ensuring your images are normalized to the ImageNet mean and std-dev is critical for the model to recognize high-level features like textures and edges during the transfer learning process.
# Import the os module for filesystem path and directory operations import os # Import shutil to copy files and manage file operations import shutil # Import random to shuffle image lists before splitting into sets import random # Define a helper function to create Train, Val, and Test folders for each category def create_folders ( base_path , categories ): # Loop over every detected category name for category in categories : # Create the Train subfolder for the current category (if it doesn't already exist) os . makedirs ( os . path . join ( base_path , ' Train ' , category ), exist_ok =True ) # Create the Val subfolder for the current category (if it doesn't already exist) os . makedirs ( os . path . join ( base_path , ' Val ' , category ), exist_ok =True ) # Create the Test subfolder for the current category (if it doesn't already exist) os . makedirs ( os . path . join ( base_path , ' Test ' , category ), exist_ok =True ) # Define a function to split data into train, validation, and test subsets def split_data ( source_folder , dest_folder , train_ratio = 0.7 , validate_ratio = 0.2 ): # Get the list of subfolders (categories) inside the source folder categories = [ d for d in os . listdir ( source_folder ) if os . path . isdir ( os . path . join ( source_folder , d ))] # Ensure all required destination folders (Train/Val/Test per category) exist create_folders ( dest_folder , categories ) # Iterate over each category to process its images for category in categories : # Build the full path to the current category directory category_path = os . path . join ( source_folder , category ) # List all image files within this category directory images = [ f for f in os . listdir ( category_path ) if os . path . isfile ( os . path . join ( category_path , f ))] # Randomly shuffle the images to avoid ordering bias random . shuffle ( images ) # Calculate the index at which the training subset ends train_split = int ( len ( images ) * train_ratio ) # Calculate the index at which the validation subset ends validate_split = int ( len ( images ) * ( train_ratio + validate_ratio )) # Select the training images based on the first split train_images = images [: train_split ] # Select the validation images between train_split and validate_split validate_images = images [ train_split : validate_split ] # The remaining images belong to the test set test_images = images [ validate_split :] # Copy each training image into the corresponding Train/category folder for image in train_images : shutil . copy ( os . path . join ( category_path , image ), os . path . join ( dest_folder , ' Train ' , category , image )) # Copy each validation image into the corresponding Val/category folder for image in validate_images : shutil . copy ( os . path . join ( category_path , image ), os . path . join ( dest_folder , ' Val ' , category , image )) # Copy each test image into the corresponding Test/category folder for image in test_images : shutil . copy ( os . path . join ( category_path , image ), os . path . join ( dest_folder , ' Test ' , category , image )) # Define the original dataset folder containing the class subfolders source_folder = ' D:/Data-Sets-Image-Classification/Star-Wars-Characters ' # Define the destination folder where the Train/Val/Test structure will be created dest_folder = ' D:/Data-Sets-Image-Classification/Star-Wars-Characters-For-Classification ' # Call the split_data function to perform the splitting operation split_data ( source_folder , dest_folder ) With your images split into folders, the model can now iterate over them in training and validation loops.
When you train FasterViT PyTorch custom dataset scripts, the Normalize transform is non-negotiable. FasterViT expects specific mean and standard deviation values (usually [0.485, 0.456, 0.406]) to align your custom pixel data with its pre-trained ImageNet weights. Failing to match these will cause the Hierarchical Attention (HAT) modules to converge much slower, leading to poor accuracy even after many epochs.
Configuring the FasterViT Backbone: Understanding HAT and Carrier Tokens Now comes the heart of the pipeline: training the FasterViT model on your custom dataset.
This code uses standard PyTorch structures like data loaders, optimizers, schedulers, and loss functions.
The core innovation of FasterViT is its Hierarchical Attention (HAT) . Unlike standard ViTs that compute attention across every pixel patch (which is computationally expensive), FasterViT breaks the image into local windows and uses ‘carrier tokens’ to summarize information across those windows. This allows the model to capture global context with a fraction of the memory footprint, making it ideal for high-resolution classification tasks.
To effectively train FasterViT PyTorch custom dataset models, use a slightly lower learning rate (e.g., 1e-5) than you would for a standard ResNet. This is a critical step when you train FasterViT PyTorch custom dataset logic because it prevents the Transformer blocks from ‘forgetting’ complex spatial relationships. Always monitor your validation loss to ensure your train FasterViT PyTorch custom dataset process isn’t overfitting on small data samples.
# Import os to handle file paths and directory operations import os # Import the core PyTorch library import torch # Import torchvision datasets and transforms for image loading and preprocessing from torchvision import datasets , transforms # Import DataLoader to batch and iterate over datasets from torch . utils . data import DataLoader # Import create_model from fastervit to construct the FasterViT architecture from fastervit import create_model # Import PyTorch's optimization module import torch . optim as optim # Import learning rate scheduler utilities from torch . optim import lr_scheduler # Import time to measure training duration import time # Import copy to deep copy model weights when tracking the best model import copy # Define a training loop function for the model def train_model ( model , criterion , optimizer , scheduler , num_epochs ): # Record the starting time of the training process since = time . time () # Make a deep copy of the model's initial weights to store the best version best_model_wts = copy . deepcopy ( model . state_dict ()) # Initialize the best accuracy with zero best_acc = 0.0 # Loop over each epoch in the training process for epoch in range ( num_epochs ): # Print the current epoch index and total epochs print ( f 'Epoch { epoch } / { num_epochs - 1 } ' ) # Print a visual separator line print ( ' - ' * 10 ) # Each epoch has both a training phase and a validation phase for phase in [ ' train ' , ' val ' ]: # Set the model to training mode during the train phase if phase == ' train ' : model . train () # Set the model to evaluation mode during the validation phase else : model . eval () # Initialize running loss for the epoch running_loss = 0.0 # Initialize running correct predictions for the epoch running_corrects = 0 # Iterate over batches from the dataloader of the current phase for inputs , labels in dataloaders [ phase ]: # Move input images to the selected device (CPU or GPU) inputs = inputs . to ( device ) # Move labels to the selected device labels = labels . to ( device ) # Reset gradients of the optimizer at the start of each batch optimizer . zero_grad () # Enable gradient computation only when in training phase with torch . set_grad_enabled ( phase == ' train ' ): # Perform a forward pass through the model to get outputs outputs = model ( inputs ) # Get the predicted class indices by taking the max logit _ , preds = torch . max ( outputs , 1 ) # Compute the loss between model outputs and true labels loss = criterion ( outputs , labels ) # If in training phase, perform backpropagation and optimizer step if phase == ' train ' : # Backpropagate the loss loss . backward () # Update model parameters optimizer . step () # Accumulate the batch loss scaled by the batch size running_loss += loss . item () * inputs . size ( 0 ) # Accumulate the number of correct predictions running_corrects += torch . sum ( preds == labels . data ) # Step the learning rate scheduler after finishing the training phase if phase == ' train ' : scheduler . step () # Compute the epoch loss by dividing total loss by dataset size epoch_loss = running_loss / dataset_sizes [ phase ] # Compute the epoch accuracy as corrects divided by dataset size epoch_acc = running_corrects . double () / dataset_sizes [ phase ] # Print the loss and accuracy for this phase print ( f ' { phase } Loss: { epoch_loss :.4f } Acc: { epoch_acc :.4f } ' ) # If we are in validation phase and this accuracy is the best so far if phase == ' val ' and epoch_acc > best_acc : # Update the best accuracy value best_acc = epoch_acc # Save the current model weights as the best model best_model_wts = copy . deepcopy ( model . state_dict ()) # Print a blank line for better readability between epochs print () # Compute total training time in seconds time_elapsed = time . time () - since # Print the total training time in minutes and seconds print ( f 'Training complete in { time_elapsed // 60 :.0f } m { time_elapsed % 60 :.0f } s' ) # Print the best validation accuracy achieved during training print ( f 'Best val Acc: { best_acc :4f } ' ) # Load the best model weights back into the model model . load_state_dict ( best_model_wts ) # Return the best model after training return model # Use the main guard to ensure code only runs when this script is executed directly if __name__ == " __main__ " : # Set the path to the prepared Train/Val dataset directory data_dir = " D:/Data-Sets-Image-Classification/Star-Wars-Characters-For-Classification " # Define image transformations for training and validation datasets datatrasforms = { # Training data augmentation and normalization pipeline ' train ' : transforms . Compose ([ # Resize the shortest side of the image to 256 pixels transforms . Resize ( 256 ), # Randomly crop a 224x224 patch from the image transforms . RandomResizedCrop ( 224 ), # Randomly flip the image horizontally for augmentation transforms . RandomHorizontalFlip (), # Convert the image to a PyTorch tensor transforms . ToTensor (), # Normalize the image with ImageNet mean and std transforms . Normalize ([ 0.485 , 0.456 , 0.406 ], [ 0.229 , 0.224 , 0.225 ]) ]), # Validation data preprocessing and normalization pipeline ' val ' : transforms . Compose ([ # Resize the shortest side of the image to 256 pixels transforms . Resize ( 256 ), # Take a centered 224x224 crop from the image transforms . CenterCrop ( 224 ), # Convert the image to a PyTorch tensor transforms . ToTensor (), # Normalize the image with ImageNet mean and std transforms . Normalize ([ 0.485 , 0.456 , 0.406 ], [ 0.229 , 0.224 , 0.225 ]) ]), } # Create ImageFolder datasets for train and val from the directory structure image_datasets = { x : datasets . ImageFolder ( os . path . join ( data_dir , x ), datatrasforms [ x ]) for x in [ ' train ' , ' val ' ]} # Wrap datasets with DataLoaders for batching and shuffling dataloaders = { x : DataLoader ( image_datasets [ x ], batch_size = 32 , shuffle =True , num_workers = 4 ) for x in [ ' train ' , ' val ' ]} # Get dataset sizes for calculating loss and accuracy dataset_sizes = { x : len ( image_datasets [ x ]) for x in [ ' train ' , ' val ' ]} # Extract class names from the training dataset folder structure class_names = image_datasets [ ' train ' ]. classes # Choose GPU if available; otherwise, fall back to CPU device = torch . device ( " cuda:0 " if torch . cuda . is_available () else " cpu " ) # Create a FasterViT model instance with pretrained weights loaded from given path model = create_model ( ' faster_vit_0_224 ' , pretrained =True , model_path = " d:/temp/models/faster_vit_0.pth.tar " ) # Get the number of features from the model head num_ftrs = model . head . in_features # Replace the final classification layer with a new Linear layer for our number of classes model . head = torch . nn . Linear ( num_ftrs , len ( class_names )) # Move the model to the selected device (GPU or CPU) model = model . to ( device ) # Define the cross-entropy loss function for multi-class classification criterion = torch . nn . CrossEntropyLoss () # Create an SGD optimizer with a learning rate and momentum optimizer = optim . SGD ( model . parameters (), lr = 0.001 , momentum = 0.9 ) # Setup a StepLR scheduler to reduce LR every 7 epochs by a factor of 0.1 scheduler = lr_scheduler . StepLR ( optimizer , step_size = 7 , gamma = 0.1 ) # Train the model using the train_model function for 100 epochs model = train_model ( model , criterion , optimizer , scheduler , num_epochs = 100 ) # Save the trained model weights to disk torch . save ( model . state_dict (), ' d:/temp/models/star_wars_faster_vit_model.pth ' ) Training on your custom dataset ensures the model learns distinct visual differences between your classes.
We utilize the AdamW optimizer here because Transformers are notoriously sensitive to weight decay. Standard SGD often fails to capture the complex relationships in attention heads. By decoupling the weight decay from the gradient update, we ensure the model learns specific character features—like Vader’s helmet or Yoda’s skin tone—without over-parameterizing the weights and causing the model to diverge.
Monitoring the loss curve is essential when training FasterViT Image Classification with a Custom Dataset in Python . Since FasterViT uses a hierarchical approach, the initial epochs might show slower convergence compared to traditional CNNs, but the final accuracy is typically superior. By the end of this FasterViT Image Classification with a Custom Dataset in Python tutorial, you will see that the AdamW optimizer helps the transformer heads focus on the most discriminative parts of the Star Wars character images.
Testing the Classifier: Real-World Inference on Star Wars Data Once training completes, you want to verify that the model works on unseen data.
Putting the predicted label onto the image makes it easy to visually confirm the model’s performance.
o verify you managed to train FasterViT PyTorch custom dataset code correctly, measure the ‘Images Per Second’ (throughput). You’ll find that a well-executed train FasterViT PyTorch custom dataset pipeline allows for nearly 2x faster processing than standard ViTs. This performance gain is the primary reason developers choose to train FasterViT PyTorch custom dataset models for real-time edge AI applications.
# Import PyTorch for model operations and tensors import torch # Import transforms for image preprocessing steps from torchvision import transforms # Import the FasterViT model creation utility from fastervit import create_model # Import os to work with filesystem paths import os # Import OpenCV for image reading and display import cv2 # Import NumPy for array operations import numpy as np # Set the initial number of classes (will be updated later) num_classes = 50 # Create a FasterViT model instance with the specified configuration model = create_model ( ' faster_vit_0_224 ' , pretrained =False ) # Adjust the classification head to match the desired number of classes model . head = torch . nn . Linear ( model . head . in_features , num_classes ) # Select GPU if available; otherwise, use CPU device = torch . device ( " cuda:0 " if torch . cuda . is_available () else " cpu " ) # Move the model to the chosen device model = model . to ( device ) # Define the path to the file containing the saved model weights model_path = ' d:/temp/models/star_wars_faster_vit_model.pth ' # Load the saved model weights from disk into the model model . load_state_dict ( torch . load ( model_path , map_location = device )) # Put the model into evaluation mode to disable dropout and other training layers model . eval () # set the model to evaluation mode # Define preprocessing steps for input images before feeding them into the model preprocess = transforms . Compose ([ # Convert the input NumPy array to a PIL Image transforms . ToPILImage (), # Convert the numpy array to PIL Image # Resize the image to 256 pixels on the shortest side transforms . Resize (( 256 )), # Center crop the image to 224x224 pixels transforms . CenterCrop (( 224 )), # Convert the image to a PyTorch tensor transforms . ToTensor (), # Normalize the image using ImageNet mean and standard deviation transforms . Normalize ( mean = [ 0.485 , 0.456 , 0.406 ], std = [ 0.229 , 0.224 , 0.225 ]) ]) # Define a helper function to load and preprocess a single image def load_image ( image_path ): # Read the image from disk using OpenCV (BGR format) image = cv2 . imread ( image_path ) # load the image using OpenCV # Convert the image from BGR color space to RGB image = cv2 . cvtColor ( image , cv2 . COLOR_BGR2RGB ) # convert BGR to RGB # Apply the preprocessing pipeline defined above image = preprocess ( image ) # apply the transformations # Add a batch dimension to create a 4D tensor image = image . unsqueeze ( 0 ) # Move the image tensor to the selected device (GPU or CPU) image = image . to ( device ) # move the image to GPU if available # Return the preprocessed image tensor return image # Define a function that loads an image, runs it through the model, and returns the predicted class name def predict ( image_path , model , class_names ): # Load and preprocess the input image image = load_image ( image_path ) # Disable gradient computation for inference with torch . no_grad (): # Forward pass through the model to get class scores outputs = model ( image ) # Get the index of the class with the highest score _ , preds = torch . max ( outputs , 1 ) # Map the predicted index to the corresponding class name preducted_class = class_names [ preds . item ()] # Return the predicted class label return preducted_class # Import glob to list files using wildcard patterns from glob import glob # Define the path to the Test folder for retrieving class names # path for test images - to get classes names from the folder names testPath = " D:/Data-Sets-Image-Classification/Star-Wars-Characters-For-Classification/Test " # Get the list of class names by reading subfolder names inside the Test directory # Get the subfolder names (class names) from the test folder class_names = [ f for f in os . listdir ( testPath ) if os . path . isdir ( os . path . join ( testPath , f ))] # Print the detected class names to verify them print ( class_names ) # Update the number of classes based on detected folder names # define the number of classes num_classes = len ( class_names ) # Define the path to a sample image used for testing the model imagePath = " Visual-Language-Models-Tutorials/FasterViT - StarWars - Image classification on your Custom Dataset using Fast Vision Transformers/Yoda-Test-Image.jpg " # Call the predict function to obtain a predicted class label predicted_class = predict ( imagePath , model , class_names ) # Print the predicted class for inspection print ( f "Predicted class : { predicted_class } " ) # Define a function that predicts the class and draws the predicted label on the image def predict_and_draw ( image_path , model , class_names ): # Load the original image using OpenCV # load the image image = cv2 . imread ( image_path ) # load the image using OpenCV # Convert the loaded image from BGR to RGB for preprocessing image_rgb = cv2 . cvtColor ( image , cv2 . COLOR_BGR2RGB ) # convert BGR to RGB # Apply the same preprocessing used for training/validation input_tensor = preprocess ( image_rgb ) # apply the transformations # Add a batch dimension to the input tensor input_tensor = input_tensor . unsqueeze ( 0 ) # add batch dimension # Move the input tensor to the selected device input_tensor = input_tensor . to ( device ) # move the image to GPU if available # Disable gradient computation during inference with torch . no_grad (): # Run the model to get output logits outputs = model ( input_tensor ) # Find the index of the highest scoring class _ , preds = torch . max ( outputs , 1 ) # Convert the index to the corresponding class name predicted_class = class_names [ preds . item ()] # Prepare the text label to overlay on the image # draw the label on the image text = f "Predicted: { predicted_class } " # Choose the font face for the text font = cv2 . FONT_HERSHEY_SIMPLEX # Set font scale (size of the text) font_scale = 1 # Set the thickness of the text stroke font_thickness = 3 # Choose the initial position (x, y) for the text on the image text_x , text_y = 10 , 50 # Draw the predicted label on the image using the specified font and color cv2 . putText ( image , text , ( text_x , text_y ), font , font_scale , ( 0 , 100 , 100 ), font_thickness ) # Display the result image in a window titled "Predicted Image" # Display the image with the label cv2 . imshow ( " Predicted Image " , image ) # Wait for a key press before closing the window cv2 . waitKey ( 0 ) # Close all OpenCV windows cv2 . destroyAllWindows () # Set the path where the labeled output image will be saved # Save the image with the label ouput_image_path = " D:/temp/predicted_image.jpg " # Write the modified image with the prediction to disk cv2 . imwrite ( ouput_image_path , image ) # Print the path to the saved predicted image print ( f "Predicted image saved at: { ouput_image_path } " ) # Call the helper function to predict and visualize the class on the test image # Run the function on a test image predict_and_draw ( imagePath , model , class_names ) Testing confirms that the model can generalize to new images and gives you visual feedback on its performance.
FAQ What is FasterViT image classification? FasterViT image classification uses a hybrid model combining convolution and transformer layers to learn image features and assign class labels efficiently.
Why split the dataset into train, val, and test? Splitting allows the model to learn patterns (train), tune performance (val), and evaluate generalization (test) for reliable results.
What does the scheduler do in training? The scheduler reduces the learning rate over time, helping stabilize training and improve final accuracy.
Why normalize images before training? Normalization ensures images have consistent pixel statistics, which helps the model converge faster.
Do I need GPU for this tutorial? No, but GPU speeds up training significantly, especially with large datasets and transformer blocks.
Can I reuse the trained model for other datasets? Yes, you can fine-tune it or retrain the head for different classes.
What library provides FasterViT? FasterVit library offers implementations of the FasterViT architecture compatible with PyTorch.
How do I visualize predictions? Predictions are written onto images using OpenCV’s text overlay and display functions.
What is the main loss function used? CrossEntropyLoss is used, which is common for multi-class classification tasks.
Summary In this complete FasterViT image classification tutorial, you learned how to:
✔ Set up a stable Python + PyTorch environment
FasterViT’s hybrid architecture gives you the speed of CNNs and the global context power of transformers, perfect for modern image classification tasks.
Once your classification model is trained, the next step is implementation; check out How to Use FasterViT for Image and Video Classification for real-time application tips.
Conclusion You now have a complete, end-to-end FasterViT image classification workflow using a custom dataset.
Transformer-based architectures like FasterViT bring the capability to understand global image context while keeping the efficiency of convolutional representations.
Connect : ☕ Buy me a coffee — https://ko-fi.com/eranfeit
🖥️ Email : feitgemel@gmail.com
🌐 https://eranfeit.net
🤝 Fiverr : https://www.fiverr.com/s/mB3Pbb
Enjoy,
Eran
Skip to content 
