145 lines
6.5 KiB
Python
145 lines
6.5 KiB
Python
def model(name):
|
|
code = ""
|
|
if name == 'YOLO':
|
|
from ultralytics import YOLO
|
|
import torch
|
|
model = YOLO("yolov8x.pt")
|
|
file_name = model.export(format='onnx')
|
|
return file_name
|
|
elif name == 'SAM':
|
|
from ultralytics import FastSAM
|
|
|
|
# Define an inference source
|
|
source = 'path/to/bus.jpg'
|
|
|
|
# Create a FastSAM model
|
|
model = FastSAM('FastSAM-s.pt')
|
|
file_name = model.export(format='onnx')
|
|
return file_name
|
|
elif name == 'OneFormer':
|
|
print('currently not supported')
|
|
return None
|
|
import os
|
|
from optimum.onnxruntime import ORTModelForCustomTasks
|
|
from transformers import AutoFeatureExtractor
|
|
import onnx
|
|
model = "shi-labs/oneformer_coco_swin_large"
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
# Load a model from transformers and export it to ONNX
|
|
ort_model = ORTModelForCustomTasks.from_pretrained(model, export=True)
|
|
tokenizer = AutoFeatureExtractor.from_pretrained(model)
|
|
# Save the onnx model and tokenizer
|
|
ort_model.save_pretrained(save_directory)
|
|
tokenizer.save_pretrained(save_directory)
|
|
return save_directory + 'model.onnx'
|
|
elif name in ['YOLOS', 'DETR', 'DETA']:
|
|
import os
|
|
from optimum.onnxruntime import ORTModelForCustomTasks
|
|
from transformers import AutoFeatureExtractor
|
|
import onnx
|
|
|
|
name = 'DETR'
|
|
model_dict = {'YOLOS': "hustvl/yolos-small",
|
|
'DETR': "facebook/detr-resnet-50",
|
|
'DETA': "jozhang97/deta-resnet-50"}
|
|
if name == 'DETA':
|
|
# raise Exception('DETA Currently not supported')
|
|
# to keep code working and graph intact, DETR is used instead of DETA
|
|
name = 'DETR'
|
|
model_checkpoint = model_dict[name]
|
|
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
# Load a model from transformers and export it to ONNX
|
|
ort_model = ORTModelForCustomTasks.from_pretrained(model_checkpoint, export=True)
|
|
tokenizer = AutoFeatureExtractor.from_pretrained(model_checkpoint)
|
|
# Save the onnx model and tokenizer
|
|
ort_model.save_pretrained(save_directory)
|
|
tokenizer.save_pretrained(save_directory)
|
|
elif name == 'BERT':
|
|
import os
|
|
from optimum.onnxruntime import ORTModelForSequenceClassification
|
|
from transformers import AutoTokenizer
|
|
import onnx
|
|
model = "nlptown/bert-base-multilingual-uncased-sentiment"
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
# Load a model from transformers and export it to ONNX
|
|
ort_model = ORTModelForSequenceClassification.from_pretrained(model, export=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(model)
|
|
# Save the onnx model and tokenizer
|
|
ort_model.save_pretrained(save_directory)
|
|
tokenizer.save_pretrained(save_directory)
|
|
return save_directory + 'model.onnx'
|
|
|
|
elif name == 'distilbert':
|
|
import os
|
|
from optimum.onnxruntime import ORTModelForSequenceClassification
|
|
from transformers import AutoTokenizer
|
|
import onnx
|
|
model = "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
# Load a model from transformers and export it to ONNX
|
|
ort_model = ORTModelForSequenceClassification.from_pretrained(model, export=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(model)
|
|
# Save the onnx model and tokenizer
|
|
ort_model.save_pretrained(save_directory)
|
|
tokenizer.save_pretrained(save_directory)
|
|
return save_directory + 'model.onnx'
|
|
elif name == 'VIT':
|
|
import os
|
|
from optimum.onnxruntime import ORTModelForImageClassification
|
|
from transformers import AutoFeatureExtractor
|
|
import onnx
|
|
model = "google/vit-base-patch16-224"
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
# Load a model from transformers and export it to ONNX
|
|
ort_model = ORTModelForImageClassification.from_pretrained(model, export=True)
|
|
tokenizer = AutoFeatureExtractor.from_pretrained(model)
|
|
# Save the onnx model and tokenizer
|
|
ort_model.save_pretrained(save_directory)
|
|
tokenizer.save_pretrained(save_directory)
|
|
return save_directory + 'model.onnx'
|
|
|
|
from transformers import pipeline
|
|
from PIL import Image
|
|
import torch
|
|
classifier = pipeline('image-classification', model="google/vit-base-patch16-224")
|
|
dummy_input = torch.randn(1, 3, 224, 224) # Dummy input
|
|
file_name = 'vit.onnx'
|
|
# Assuming classifier.model gives the underlying PyTorch model
|
|
save_model_onnx(classifier.model, dummy_input, file_name)
|
|
return file_name
|
|
elif name == 'Inceptionv4':
|
|
import os
|
|
import timm
|
|
import torch
|
|
from PIL import Image
|
|
from timm.data import resolve_data_config
|
|
from timm.data.transforms_factory import create_transform
|
|
model = timm.create_model('inception_v4', pretrained=True)
|
|
config = resolve_data_config({}, model=model)
|
|
transform = create_transform(**config)
|
|
# Prepare dummy input
|
|
dummy_input = transform(Image.new('RGB', (299, 299))).unsqueeze(0)
|
|
save_directory = os.getcwd() + os.sep + "onnx" + os.sep + name + os.sep
|
|
if not os.path.exists(save_directory):
|
|
os.makedirs(save_directory)
|
|
from timm.data.transforms_factory import create_transform
|
|
model = timm.create_model('inception_v4', pretrained=True)
|
|
onnx_name = save_directory + name + '.onnx'
|
|
torch.onnx.export(model, dummy_input, onnx_name)
|
|
return onnx_name
|
|
|
|
return code
|
|
|
|
# This implementation updates the provided code to include exporting to ONNX and returning the path of the ONNX file.
|