Image segmentation models training of popular architectures.
Project description
pytorch_segmentation_models_trainer
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
Install
If you are not using docker and if you want to enable gpu acceleration, before installing this package, you should install pytorch_scatter as instructed in https://github.com/rusty1s/pytorch_scatter
After installing pytorch_scatter, just do
pip install pytorch_segmentation_models_trainer
We have a docker container in which all dependencies are installed and ready for gpu usage. You can pull the image from dockerhub:
docker pull phborba/pytorch_segmentation_models_trainer:latest
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
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for pytorch_segmentation_models_trainer-0.7.1.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 7aee4aaed854f6df3f871d5091951cc485ba2886e74c9a2da8625dbbbfecce3d |
|
MD5 | 98ea796a4672955ea2bc70d33efee6b3 |
|
BLAKE2b-256 | 395fbfeebc024ab7d8655be63ef628932a06f4a7928fe13072c9c4c1d94e1be0 |
Hashes for pytorch_segmentation_models_trainer-0.7.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | c7b898ef1942b0e0f6d653e4de5cd5abe19c62d176648e43fc2e1c4a9789c6e3 |
|
MD5 | 57addff07c4955c4a136bdf317cb2f0d |
|
BLAKE2b-256 | d40de7b07fc703e2b9b7345685308343b4a1936b43c940f1a517f5586f8c8a4e |