mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-23 04:26:35 +00:00
add export to FP16
add inference with possibility to have different
This commit is contained in:
+5
-14
@@ -31,22 +31,13 @@ def export_onnx(model_path):
|
||||
return Path(model_path).stem + '.onnx'
|
||||
|
||||
|
||||
def export_tensorrt(model_path, dataset_yaml):
|
||||
form_data_sample(path.join(path.dirname(dataset_yaml), 'minival', 'images'))
|
||||
model = YOLO(model_path)
|
||||
with open(dataset_yaml, 'r') as file:
|
||||
yaml_data = yaml.safe_load(file) or {}
|
||||
yaml_data['minival'] = 'minival/images'
|
||||
with open(dataset_yaml, 'w') as file:
|
||||
yaml.dump(yaml_data, file)
|
||||
|
||||
model.export(
|
||||
def export_tensorrt(model_path):
|
||||
YOLO(model_path).export(
|
||||
format='engine',
|
||||
batch=4,
|
||||
half=True,
|
||||
nms=True,
|
||||
data=dataset_yaml,
|
||||
split='minival'
|
||||
simplify=True,
|
||||
nms=True
|
||||
)
|
||||
|
||||
def form_data_sample(destination_path, size=500, write_txt_log=False):
|
||||
@@ -76,6 +67,6 @@ def show_model(model: str = None):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
export_tensorrt('azaion-2025-03-10.pt', path.join(datasets_dir, 'azaion-2025-03-10', 'data.yaml'))
|
||||
export_tensorrt('azaion-2025-03-10.pt')
|
||||
# export_rknn('azaion-2025-03-10.pt')
|
||||
# export_onnx('azaion-2025-03-10.pt')
|
||||
Reference in New Issue
Block a user