Pytorch utilities for model training on GPU and TPU in a single, flexible interface that can be subclassed with your own methods. Model results are (optionally) saved to a MongoDB, for asynchronous visualization.
To install run:
git clone https://github.com/anayebi/ptutils
cd ptutils/
pip install -e .
The example scripts support training ResNet-18 on ImageNet categorization, e.g.
cd ptutils/model_training/
python runner.py --config=configs/resnet18_supervised_imagenet_trainer_[gpu/tpu].json
You can substitute your own training method by importing from ptutils.model_training.runner import Runner
, and subclassing Runner.train()
.
By default, this packages saves model results to MongoDB.
If you would like to use it, follow these instructions to install MongoDB.
Otherwise, to disable this feature, set "use_mongodb": false
in your configuration json.
The function ptutils.core.utils.grab_results()
is an example of how to grab the results from MongoDB for the SupervisedImageNetTrainer
, and this notebook gives an example of plotting it.
Put this in .git/hooks/pre-commit
, and run sudo chmod +x .git/hooks/pre-commit
.
#!/usr/bin/env bash
echo "# Running pre-commit hook"
echo "#########################"
echo "Checking formatting"
format_occurred=false
declare -a black_dirs=("ptutils/" "setup.py")
for black_dir in "${black_dirs[@]}"; do
echo ">>> Checking $black_dir"
black --check "$black_dir"
if [ $? -ne 0 ]; then
echo ">>> Reformatting now!"
black "$black_dir"
format_occurred=true
fi
done
if [ "$format_occurred" = true ]; then
exit 1
fi
MIT
- Aran Nayebi (Stanford/MIT)
- Nathan C. L. Kong (Stanford)
- Javier Sagastuy-Brena (Stanford)
If you have any questions or encounter issues, either submit a Github issue here (preferred) or email me.