pytorch-tools icon indicating copy to clipboard operation
pytorch-tools copied to clipboard

Tool box for PyTorch

Pytorch-Tools

Tool box for PyTorch for fast prototyping.

Overview

  • FitWrapper - Keras like model trainer
  • Losses - collection of different Loss functions.
  • Metrics - collection of metrics.
  • Models - classification model zoo.
  • Optimizers
  • Segmentation Models - segmentation models zoo
  • TTA wrapper - wrapper for easy test-time augmentation

Installation

Requires GPU drivers and CUDA already installed.

pip install git+https://github.com/bonlime/pytorch-tools.git@master

It is also recommended to install NVIDIA Apex to allow usage of additional optimizers

pip install ---upgrade -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" git+https://github.com/NVIDIA/apex.git