Add more model frameworks to the list

This commit is contained in:
allegroai 2022-08-20 22:57:20 +03:00
parent 5a83aa433d
commit ca384aa75c

View File

@ -40,6 +40,7 @@ class Framework(Options):
tensorflowjs = 'TensorFlow_js'
tensorflowlite = 'TensorFlow_Lite'
pytorch = 'PyTorch'
torchscript = 'TorchScript'
caffe = 'Caffe'
caffe2 = 'Caffe2'
onnx = 'ONNX'
@ -55,6 +56,8 @@ class Framework(Options):
parquet = 'Parquet'
megengine = 'MegEngine'
catboost = 'CatBoost'
tensorrt = 'TensorRT'
openvino = 'OpenVINO'
__file_extensions_mapping = {
'.pb': (tensorflow, tensorflowjs, onnx, ),
@ -84,6 +87,7 @@ class Framework(Options):
'.pkl': (scikitlearn, keras, xgboost, megengine),
'.parquet': (parquet, ),
'.cbm': (catboost, ),
'.plan': (tensorrt,),
}
__parent_mapping = {