This repository contains a Jupyter notebook for semantic segmentation using the DeepLabV3 model with a ResNet-50 backbone. The notebook demonstrates the process of loading data, training the model, and making predictions. Additionally, Ray Tune is utilized for hyperparameter tuning to optimize the model's performance.
- Python 3.6 or later
- PyTorch
- torchvision
- numpy
- matplotlib
- Pillow
- ray[tune]
- Clone the repository:
git clone https://github.com/thivyanth/deeplabv3-raytune.git
cd deeplabv3-raytune
- Install the required packages:
pip install torch torchvision numpy matplotlib pillow ray[tune]
- Place your training images in the
train-data/image
directory. - Place your training labels in the
train-data/label
directory. - Place your unlabeled images in the
unlabeled-data
directory. - Place your test images in the
test-data/image
directory.
The training process includes:
- Defining a custom dataset class (
ImageDataset
) for loading images and labels. - Initializing and customizing a pre-trained DeepLabV3 model.
- Training the model using the defined dataset and saving the trained model.
To train the model, run the notebook cells related to training.
Ray Tune is used to optimize hyperparameters for better model performance. The following steps are included for hyperparameter tuning:
- Define the Search Space: Specify the range of hyperparameters to explore.
- Configure the Scheduler and Search Algorithm: Use Ray Tune's built-in schedulers and search algorithms for efficient hyperparameter tuning.
- Run the Tuning: Execute the tuning process and log the results.
- Load the trained model weights.
- Create a test dataset and DataLoader.
- Perform inference and save predictions as
.npy
files. - Zip the predictions for easy download.
The training loss is plotted to visualize the training progress. The model predictions are saved and zipped for further analysis.
- PyTorch
- Ray Tune