mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-23 00:56:35 +00:00
add download big engine file to cdn manager
revise onnx export process fixes
This commit is contained in:
+27
-11
@@ -1,13 +1,15 @@
|
||||
import os
|
||||
import shutil
|
||||
from os import path, scandir, makedirs
|
||||
from pathlib import Path
|
||||
import random
|
||||
|
||||
import netron
|
||||
import yaml
|
||||
from ultralytics import YOLO
|
||||
|
||||
import constants
|
||||
from azaion_api import Api, ApiCredentials
|
||||
from api_client import ApiClient, ApiCredentials
|
||||
from cdn_manager import CDNManager, CDNCredentials
|
||||
from constants import datasets_dir, processed_images_dir
|
||||
from security import Security
|
||||
@@ -24,15 +26,21 @@ def export_rknn(model_path):
|
||||
pass
|
||||
|
||||
|
||||
def export_onnx(model_path):
|
||||
def export_onnx(model_path, batch_size=4):
|
||||
model = YOLO(model_path)
|
||||
onnx_path = Path(model_path).stem + '.onnx'
|
||||
if path.exists(onnx_path):
|
||||
os.remove(onnx_path)
|
||||
|
||||
model.export(
|
||||
format="onnx",
|
||||
imgsz=1280,
|
||||
batch=2,
|
||||
batch=batch_size,
|
||||
simplify=True,
|
||||
nms=True)
|
||||
return Path(model_path).stem + '.onnx'
|
||||
nms=True,
|
||||
device=0
|
||||
)
|
||||
return onnx_path
|
||||
|
||||
|
||||
def export_tensorrt(model_path):
|
||||
@@ -79,17 +87,25 @@ def upload_model(model_path: str, filename: str, size_small_in_kb: int=3):
|
||||
key = Security.get_model_encryption_key()
|
||||
model_encrypted = Security.encrypt_to(model_bytes, key)
|
||||
|
||||
part1_size = min(size_small_in_kb * 1024, int(0.9 * len(model_encrypted)))
|
||||
part1_size = min(size_small_in_kb * 1024, int(0.3 * len(model_encrypted)))
|
||||
model_part_small = model_encrypted[:part1_size] # slice bytes for part1
|
||||
model_part_big = model_encrypted[part1_size:]
|
||||
|
||||
with open(constants.CONFIG_FILE, "r") as f:
|
||||
config_dict = yaml.safe_load(f)
|
||||
d_config = Dotdict(config_dict)
|
||||
cdn_c = Dotdict(d_config.cdn)
|
||||
api_c = Dotdict(d_config.api)
|
||||
cdn_manager = CDNManager(CDNCredentials(cdn_c.host, cdn_c.access_key, cdn_c.secret_key))
|
||||
cdn_manager.upload(cdn_c.bucket, f'{filename}.big', model_part_big)
|
||||
api = ApiClient(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
|
||||
|
||||
api = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
|
||||
api.upload_file(f'{filename}.small', model_part_small)
|
||||
yaml_bytes = api.load_bytes(constants.CDN_CONFIG, '')
|
||||
data = yaml.safe_load(yaml_bytes)
|
||||
creds = CDNCredentials(data["host"],
|
||||
data["downloader_access_key"],
|
||||
data["downloader_access_secret"],
|
||||
data["uploader_access_key"],
|
||||
data["uploader_access_secret"])
|
||||
|
||||
cdn_manager = CDNManager(creds)
|
||||
|
||||
api.upload_file(f'{filename}.small', model_part_small, constants.MODELS_FOLDER)
|
||||
cdn_manager.upload(constants.MODELS_FOLDER, f'{filename}.big', model_part_big)
|
||||
|
||||
Reference in New Issue
Block a user