Skip to main content

Image segmentation models training of popular architectures.

Project description

pytorch_segmentation_models_trainer

Torch Pytorch Lightning Hydra Segmentation Models Python application Upload Python Package PyPI Publish Docker image maintainer DOI

Framework based on Pytorch, Pytorch Lightning, segmentation_models.pytorch and hydra to train semantic segmentation models using yaml config files as follows:

model:
  _target_: segmentation_models_pytorch.Unet
  encoder_name: resnet34
  encoder_weights: imagenet
  in_channels: 3
  classes: 1

loss:
  _target_: segmentation_models_pytorch.utils.losses.DiceLoss

optimizer:
  _target_: torch.optim.AdamW
  lr: 0.001
  weight_decay: 1e-4

hyperparameters:
  batch_size: 1
  epochs: 2
  max_lr: 0.1

pl_trainer:
  max_epochs: ${hyperparameters.batch_size}
  gpus: 0

train_dataset:
  _target_: pytorch_segmentation_models_trainer.dataset_loader.dataset.SegmentationDataset
  input_csv_path: /path/to/input.csv
  data_loader:
    shuffle: True
    num_workers: 1
    pin_memory: True
    drop_last: True
    prefetch_factor: 1
  augmentation_list:
    - _target_: albumentations.HueSaturationValue
      always_apply: false
      hue_shift_limit: 0.2
      p: 0.5
    - _target_: albumentations.RandomBrightnessContrast
      brightness_limit: 0.2
      contrast_limit: 0.2
      p: 0.5
    - _target_: albumentations.RandomCrop
      always_apply: true
      height: 256
      width: 256
      p: 1.0
    - _target_: albumentations.Flip
      always_apply: true
    - _target_: albumentations.Normalize
      p: 1.0
    - _target_: albumentations.pytorch.transforms.ToTensorV2
      always_apply: true

val_dataset:
  _target_: pytorch_segmentation_models_trainer.dataset_loader.dataset.SegmentationDataset
  input_csv_path: /path/to/input.csv
  data_loader:
    shuffle: True
    num_workers: 1
    pin_memory: True
    drop_last: True
    prefetch_factor: 1
  augmentation_list:
    - _target_: albumentations.Resize
      always_apply: true
      height: 256
      width: 256
      p: 1.0
    - _target_: albumentations.Normalize
      p: 1.0
    - _target_: albumentations.pytorch.transforms.ToTensorV2
      always_apply: true

To train a model with configuration path /path/to/config/folder and name test.yaml:

pytorch-smt --config-dir /path/to/config/folder --config-name test +mode=train

The mode can be stored in configuration yaml as well. In this case, do not pass the +mode= argument. If the mode is stored in the yaml and you want to overwrite the value, do not use the + clause, just mode= .

This module suports hydra features such as configuration composition. For further information, please visit https://hydra.cc/docs/intro

Citing:


@software{philipe_borba_2021_4574256,
  author       = {Philipe Borba},
  title        = {{phborba/pytorch\_segmentation\_models\_trainer: 
                   Version 0.1.2}},
  month        = mar,
  year         = 2021,
  publisher    = {Zenodo},
  version      = {v0.1.2},
  doi          = {10.5281/zenodo.4574256},
  url          = {https://doi.org/10.5281/zenodo.4574256}
}


Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for pytorch-segmentation-models-trainer, version 0.5.0
Filename, size File type Python version Upload date Hashes
Filename, size pytorch_segmentation_models_trainer-0.5.0-py3-none-any.whl (89.1 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size pytorch_segmentation_models_trainer-0.5.0.tar.gz (64.2 kB) File type Source Python version None Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page