Skip to content

simple pytorch segmentation tutorial with apex mixed precision and distributed training

License

Notifications You must be signed in to change notification settings

dtrimina/pytorch_segmentation_tutorial

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

58 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

simple pytorch segmentation tutorial with apex mixed precision and distributed training

version

  • python 3.7.5
  • torch 1.3.0
  • torchvision 0.4.1
  • apex 0.1

database

  • CamVid : It is a automotive dataset which contains 367 training, 101 validation, and 233 testing images. There are eleven different classes such as building, tree, sky, car, road, etc. while the twelfth class contains unlabeled data, which we ignore during training. The original frame resolution for this dataset is 960 × 720. We use 480 x 360 images in SegNet-Tutorial.
  • SUNRGBD : The SUN RGB-D dataset has 10,335 densely annotated RGB-D images taken from 20 different scenes. It also include all images data from NYU Depth v2 dataset, and selected images data from Berkeley B3DO and SUN3D dataset. Each pixel in the RGB-D images is assigned a semantic label in one of the 37 classes or the ‘unknown’ class. It has 5285 training/validation instances and 5050 testing instances. If you want to train your segmentation model with depth information.
  • Cityscapes : This dataset consists of 5000 fineannotated images, out of which 2975 are available for training, 500 for validation, and the remaining 1525 have been selected as test set. It contains 19 classes and an unlabel class. The raw resolution is 2048*1024.
  • ADE20k: ADE20K dataset is a challenge scence parsing dataset providing 150 class dense labels, which consists of 20K/2K/3K images for training/validation/test.

models

  • training on 4 taitanxp GPUs
  • miou on validation set, use model trained after the last epoch to evaluate
  • msc miou means using multi-scale images (0.5 0.75 1.0 1.25 1.5 1.75) and their flip version for evaluation
model miou msc miou params size(fp32)
unet - - 51.14MB
segnet - - 117MB
LinkNet - - 44.07MB
FC-DenseNet103 - - 35.58MB
ENet - - 1.34MB
DRN-C-26 0.671 0.702 78.67MB
Deeplabv3+_resnet50 0.746 0.765 153.94MB
Deeplabv3+_resnet101 - - 226.39MB

default training config

  • data augmentation: colorjit + randomhflip + randomscale + randomcrop
  • input image normalize: ToTensor + Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  • loss function: CrossEntropyLoss + class_weight
  • 150 epoch, Adam optimizer, initial_lr=5e-4

train and evaluate

# train on 4 GPUs
python -m torch.distributed.launch --nproc_per_node=4 train.py --config configs/cityscape_drn_c_26.json

# evaluate
python evaluate.py --logdir [run logdir] [-s] 

# Moreover, you can add [your configs].json in run_tasks.sh
sh run_tasks.sh

reference

About

simple pytorch segmentation tutorial with apex mixed precision and distributed training

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published