deep-learning-pytorch-demo

Train, test, export demo is based on Pytorch. The models include MLP, LeNet, AlexNet, vgg, GoogleNet, ResNet. The supported model types include torchscript, onnx, openvino, engine.

install

git clone https://github.com/taifyang/deep-learning-pytorch-demo  # clone
cd deep-learning-pytorch-demo
pip install -r requirements.txt  # install

Data Preparation

Download dataset mnist here and save in ./Dataset/mnist_png/.

simple usage

## train 
python train.py --models lenet

## test
python test.py --models lenet

## export
python export.py --models lenet --type onnx

Reference By

d2l-zh-pytorch