This repository contains the source code of segmentation branch of Y-Net.
This repository is organized as:
- data This directory should contain the training and validation data.
- pretrained_models_st1 This directory contains the pre-trained models. We only provide the models for Y-Net with ESP as encoding blocks and PSP as decoding blocks.
- Python files - These files contain the source code that we used to train the data
You can start training the model using below command:
python main.py
Please see the command line arguments for more details.
Note 1: Currently, we support only single GPU training. If you want to train the model on multiple-GPUs, you can use nn.DataParallel api provided by PyTorch.
Note 2: To train on a specific GPU (single), you can specify the GPU_ID using the CUDA_VISIBLE_DEVICES as:
CUDA_VISIBLE_DEVICES=2 python main.py
This will run the training program on GPU with ID 2.