- Notifications
You must be signed in to change notification settings - Fork2
convert a pytorch model to a model for edge device
License
NotificationsYou must be signed in to change notification settings
kuroko1t/nne
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
convert pytorch model for Edge Device
contents
python -m pip install -e.- edgetpu
If you want to compile pytorch model for edgetpu,install edgetpu_compiler
example compile pytorch model for edge device. Seeexample for details
comvert to onnx model
importnneimporttorchvisionimporttorchimportnumpyasnpinput_shape= (1,3,64,64)onnx_file='resnet.onnx'model=torchvision.models.resnet34(pretrained=True).cuda()nne.cv2onnx(model,input_shape,onnx_file)
comvert to tflite model
importtorchvisionimporttorchimportnumpyasnpimportnneinput_shape= (10,3,224,224)model=torchvision.models.mobilenet_v2(pretrained=True).cuda()tflite_file='mobilenet.tflite'nne.cv2tflite(model,input_shape,tflite_file)
comvert to tflite model(edge tpu)
importtorchvisionimporttorchimportnumpyasnpimportnneinput_shape= (10,3,112,112)model=torchvision.models.mobilenet_v2(pretrained=True)tflite_file='mobilenet.tflite'nne.cv2tflite(model,input_shape,tflite_file,edgetpu=True)
convert to TensorRT model
importnneimporttorchvisionimporttorchimportnumpyasnpinput_shape= (1,3,224,224)trt_file='alexnet_trt.pth'model=torchvision.models.alexnet(pretrained=True).cuda()nne.cv2trt(model,input_shape,trt_file)
- show summary model info
- dump detailed model information(node name, attrs) to json file.
- convert onnx model to tflite, simplifier
$nne -husage: nne [-h] [-a ANALYZE_PATH] [-s SIMPLYFY_PATH] [-t TFLITE_PATH] model_pathNeural Network Graph Analyzerpositional arguments: model_path model pathfor analyzingoptional arguments: -h, --help show thishelp message andexit -a ANALYZE_PATH, --analyze_path ANALYZE_PATH Specify the path to output the Node information of the modelin json format. -s SIMPLYFY_PATH, --simplyfy_path SIMPLYFY_PATH onnx model to simplyfier -t TFLITE_PATH, --tflite_path TFLITE_PATH onnx model to tflite
| format | support |
|---|---|
| tflite | ✅ |
| edge tpu | trial |
| onnx | ✅ |
| tensorRT | ✅ |
Apache 2.0