The toolkit for fast Deep Learning experiments in Computer Vision
TorchOk is based on PyTorch and utilizes PyTorch Lightning for training pipeline routines.
The toolkit consists of:
- Neural Network models which are proved to be the best not only on PapersWithCode but in practice. All models are under plug&play interface that easily connects backbones, necks and heads for reuse across tasks
- Out-of-the-box support of common Computer Vision tasks: classification, segmentation, image representation and detection coming soon
- Commonly used datasets, image augmentations and transformations (from Albumentations)
- Fast implementations of retrieval metrics (with the help of FAISS and ranx) and lots of other metrics from torchmetrics
- Export models to ONNX and the ability to test the exported model without changing the datasets
- All components can be customized by inheriting the unified interfaces: Lightning's training loop, tasks, models, datasets, augmentations and transformations, metrics, loss functions, optimizers and LR schedulers
- Training, validation and testing configurations are represented by YAML config files and managed by Hydra
- Only straightforward training techniques are implemented. No whistles and bells
Installation via pip can be done in two steps:
- Install PyTorch that meets your hardware requirements via official instructions
- Install TorchOk by running
pip install --upgrade torchok
To remove the previous installation of TorchOk environment, run:
conda remove --name torchok --allTo install TorchOk locally, run:
conda env create -f environment.ymlThis will create a new conda environment torchok with all dependencies.
Another way to install TorchOk is through Docker. The built image supports SSH access, Jupyter Lab and Tensorboard ports exposing. If you don't need any of this, just omit the corresponding arguments. Build the image and run the container:
docker build -t torchok --build-arg SSH_PUBLIC_KEY="<public key>" .
docker run -d --name <username>_torchok --gpus=all -v <path/to/workdir>:/workdir -p <ssh_port>:22 -p <jupyter_port>:8888 -p <tensorboard_port>:6006 torchokThe folder examples/configs contains YAML config files with some predefined training and inference configurations.
For a training example, we can use the default configuration examples/configs/classification_cifar10.yml, where the CIFAR-10 dataset and the classification task are specified. The CIFAR-10 dataset will be automatically downloaded into your ~/.cache/torchok/data/cifar10 folder (341 MB).
To train on all available GPU devices (default config):
python -m torchok -cp ../examples/configs -cn classification_cifar10To train on all available CPU cores:
python -m torchok -cp ../examples/configs -cn classification_cifar10 trainer.accelerator='cpu'During the training you can access the training and validation logs by starting a local TensorBoard:
tensorboard --logdir ~/.cache/torchok/logs/cifar10To automatically find the initial learning rate, we use Pytorch Lightning tuner which algorithm based on Cyclical Learning Rates for Training Neural Networks the article.
python -m torchok -cp ../examples/configs -cn classification_cifar10 +mode=find_lrTODO
For the ONNX model run, we can use the examples/configs/onnx_infer.yaml.
But first we need to define the field path_to_onnx.
To test ONNX model:
python test.py -cp examples/configs -cn onnx_infer +mode=testTo predict ONNX model:
python test.py -cp examples/configs -cn onnx_infer +mode=predictpython -m unittest discover -s tests/ -p "test_*.py"- MOBY (unsupervised training)
- InstanceSegmentationTask
- YOLOR neck + head
- DETR neck + head
- ImageNet
- Cityscapes
- Pytorch Metric Learning losses
- NT-ext (for unsupervised training)
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent