azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 08:56:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

add download big engine file to cdn manager

Browse Source 
revise onnx export process
fixes
This commit is contained in:
Alex Bezdieniezhnykh
2025-04-24 16:23:39 +03:00
parent c4732cec8b
commit 26eb3b991d
9 changed files with 187 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
azaion_api.py → api_client.py
+47 -16
View File
@@ -1,26 +1,43 @@
import io
import json
from http import HTTPStatus
from os import path
import requests
import yaml
import constants
from cdn_manager import CDNCredentials, CDNManager
from hardware_service import get_hardware_info
from security import Security
class ApiCredentials:
def __init__(self, url, email, password, folder):
def __init__(self, url, email, password):
self.url = url
self.email = email
self.password = password
self.folder = folder
class Api:
def __init__(self, credentials):
class ApiClient:
def __init__(self):
self.token = None
self.credentials = credentials
with open(constants.CONFIG_FILE, "r") as f:
config_dict = yaml.safe_load(f)
api_c = config_dict["api"]
self.credentials = ApiCredentials(api_c["url"],
api_c["email"],
api_c["password"])
yaml_bytes = self.load_bytes(constants.CDN_CONFIG, '')
data = yaml.safe_load(io.BytesIO(yaml_bytes))
creds = CDNCredentials(data["host"],
data["downloader_access_key"],
data["downloader_access_secret"],
data["uploader_access_key"],
data["uploader_access_secret"])
self.cdn_manager = CDNManager(creds)
def login(self):
response = requests.post(f'{self.credentials.url}/login',
@@ -29,8 +46,7 @@ class Api:
token = response.json()["token"]
self.token = token
def upload_file(self, filename: str, file_bytes: bytearray):
folder = self.credentials.folder
def upload_file(self, filename: str, file_bytes: bytearray, folder):
if self.token is None:
self.login()
url = f"{self.credentials.url}/resources/{folder}"
@@ -43,9 +59,7 @@ class Api:
except Exception as e:
print(f"Upload fail: {e}")
def load_bytes(self, filename, folder = None):
folder = folder or self.credentials.folder
def load_bytes(self, filename, folder):
hardware = get_hardware_info()
if self.token is None:
@@ -82,13 +96,30 @@ class Api:
print(f'Downloaded file: {filename}, {len(data)} bytes')
return data
def load_resource(self, big_part, small_part):
def load_big_small_resource(self, resource_name, folder, key):
big_part = path.join(folder, f'{resource_name}.big')
small_part = f'{resource_name}.small'
with open(big_part, 'rb') as binary_file:
encrypted_bytes_big = binary_file.read()
encrypted_bytes_small = self.load_bytes(small_part)
encrypted_bytes_small = self.load_bytes(small_part, folder)
encrypted_model_bytes = encrypted_bytes_small + encrypted_bytes_big
key = Security.get_model_encryption_key()
encrypted_bytes = encrypted_bytes_small + encrypted_bytes_big
model_bytes = Security.decrypt_to(encrypted_model_bytes, key)
return model_bytes
result = Security.decrypt_to(encrypted_bytes, key)
return result
def upload_big_small_resource(self, resource, resource_name, folder, key):
big_part_name = f'{resource_name}.big'
small_part_name = f'{resource_name}.small'
resource_encrypted = Security.encrypt_to(resource, key)
part_small_size = min(constants.SMALL_SIZE_KB * 1024, int(0.2 * len(resource_encrypted)))
part_small = resource_encrypted[:part_small_size] # slice bytes for part1
part_big = resource_encrypted[part_small_size:]
self.cdn_manager.upload(constants.MODELS_FOLDER, big_part_name, part_big)
with open(path.join(folder, big_part_name), 'wb') as f:
f.write(part_big)
self.upload_file(small_part_name, part_small, constants.MODELS_FOLDER)
config.yaml
+2 -11
View File
@@ -1,16 +1,7 @@
cdn:
host: 'https://cdnapi.azaion.com/'
downloader_access_key: '8ynZ0rrMLL00GLBopklw'
downloader_access_secret: 'McNgEKhAJUxoa3t4WDDbCbhYPg4Qhe7FNQEKrtbk'
uploader_access_key: 'YhdHtKaq8DmvrYohetu6'
uploader_access_secret: 'nlOtjo1c4UWiMiJOjcIpR0aJFPitIhcwU6zFev7H'
bucket: 'models'
api:
url: 'https://api.azaion.com'
user: 'admin@azaion.com'
pw: 'Az@1on1000Odm$n'
folder: ''
email: 'uploader@azaion.com'
password: 'Az@1on_10Upl0@der'
queue:
host: '188.245.120.247'
constants.py
+4
View File
@@ -39,3 +39,7 @@ AI_ONNX_MODEL_FILE_SMALL = "azaion.onnx.small"
AI_TENSOR_MODEL_FILE_BIG = "azaion.engine.big"
AI_TENSOR_MODEL_FILE_SMALL = "azaion.engine.small"
SMALL_SIZE_KB = 3
CDN_CONFIG = 'cdn.yaml'
MODELS_FOLDER = 'models'
exports.py
+27 -11
View File
@@ -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)
hardware_service.py
-1
View File
@@ -34,7 +34,6 @@ def get_mac_address(interface="Ethernet"):
def get_hardware_info():
is_windows = os.name == 'nt'
res = subprocess.check_output("ver", shell=True).decode('utf-8')
if "Microsoft Windows" in res:
is_windows = True
inference/tensorrt_engine.py
+73 -18
View File
@@ -8,28 +8,29 @@ import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
from inference.onnx_engine import InferenceEngine
import pycuda.autoinit # required for automatically initialize CUDA, do not remove.
# required for automatically initialize CUDA, do not remove.
import pycuda.autoinit
import pynvml
# TODO: 2. Convert onnx model with 4 batch and make sure it is working
class TensorRTEngine(InferenceEngine):
def __init__(self, model_bytes: bytes, batch_size: int = 4, **kwargs):
self.batch_size = batch_size
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def __init__(self, model_bytes: bytes, **kwargs):
try:
logger = trt.Logger(trt.Logger.WARNING)
# metadata_len = struct.unpack("<I", model_bytes[:4])[0]
# try:
# self.metadata = json.loads(model_bytes[4:4 + metadata_len])
# self.class_names = self.metadata['names']
# print(f"Model metadata: {json.dumps(self.metadata, indent=2)}")
# except json.JSONDecodeError as err:
# print(f"Failed to parse metadata")
# return
# engine_data = model_bytes[4 + metadata_len:]
metadata_len = struct.unpack("<I", model_bytes[:4])[0]
try:
self.metadata = json.loads(model_bytes[4:4 + metadata_len])
self.class_names = self.metadata['names']
print(f"Model metadata: {json.dumps(self.metadata, indent=2)}")
except json.JSONDecodeError as err:
print(f"Failed to parse metadata")
return
engine_data = model_bytes[4 + metadata_len:]
runtime = trt.Runtime(logger)
self.engine = runtime.deserialize_cuda_engine(engine_data)
runtime = trt.Runtime(self.TRT_LOGGER)
self.engine = runtime.deserialize_cuda_engine(model_bytes)
if self.engine is None:
raise RuntimeError(f"Failed to load TensorRT engine!")
@@ -55,7 +56,7 @@ class TensorRTEngine(InferenceEngine):
self.output_name = self.engine.get_tensor_name(1)
engine_output_shape = tuple(self.engine.get_tensor_shape(self.output_name))
self.output_shape = [
batch_size if self.input_shape[0] == -1 else self.input_shape[0],
4 if self.input_shape[0] == -1 else self.input_shape[0], # by default, batch size is 4
300 if engine_output_shape[1] == -1 else engine_output_shape[1], # max detections number
6 if engine_output_shape[2] == -1 else engine_output_shape[2] # x1 y1 x2 y2 conf cls
]
@@ -73,7 +74,61 @@ class TensorRTEngine(InferenceEngine):
def get_batch_size(self) -> int:
return self.batch_size
# In tensorrt_engine.py, modify the run method:
@staticmethod
def get_gpu_memory_bytes(device_id=0) -> int:
total_memory = None
try:
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
total_memory = mem_info.total
except pynvml.NVMLError:
total_memory = None
finally:
try:
pynvml.nvmlShutdown()
except pynvml.NVMLError:
pass
return 2 * 1024 * 1024 * 1024 if total_memory is None else total_memory # default 2 Gb
@staticmethod
def get_engine_filename(device_id=0) -> str | None:
try:
device = cuda.Device(device_id)
sm_count = device.multiprocessor_count
cc_major, cc_minor = device.compute_capability()
return f"azaion.cc_{cc_major}.{cc_minor}_sm_{sm_count}.engine"
except Exception:
return None
@staticmethod
def convert_from_onnx(onnx_model: bytes) -> bytes | None:
workspace_bytes = int(TensorRTEngine.get_gpu_memory_bytes() * 0.9)
explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TensorRTEngine.TRT_LOGGER) as builder, \
builder.create_network(explicit_batch_flag) as network, \
trt.OnnxParser(network, TensorRTEngine.TRT_LOGGER) as parser, \
builder.create_builder_config() as config:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
if not parser.parse(onnx_model):
return None
if builder.platform_has_fast_fp16:
print('Converting to supported fp16')
config.set_flag(trt.BuilderFlag.FP16)
else:
print('Converting to supported fp32. (fp16 is not supported)')
plan = builder.build_serialized_network(network, config)
if plan is None:
print('Conversion failed.')
return None
return bytes(plan)
def run(self, input_data: np.ndarray) -> List[np.ndarray]:
try:
requirements.txt
+1 -1
View File
@@ -9,7 +9,7 @@ opencv-python
matplotlib
PyYAML
cryptography
numpy
numpy==1.26.4
requests
pyyaml
boto3
start_inference.py
+22 -20
View File
@@ -1,14 +1,23 @@
import re
import pycuda.driver as cuda
import yaml
import constants
from azaion_api import Api, ApiCredentials
from api_client import ApiClient, ApiCredentials
from cdn_manager import CDNManager, CDNCredentials
from inference.inference import Inference
from inference.onnx_engine import OnnxEngine
from inference.tensorrt_engine import TensorRTEngine
from security import Security
from utils import Dotdict
def get_engine_filename(device_id=0):
try:
device = cuda.Device(device_id)
sm_count = device.multiprocessor_count
cc_major, cc_minor = device.compute_capability()
return f"azaion.cc_{cc_major}.{cc_minor}_sm_{sm_count}.engine"
except Exception:
return None
if __name__ == "__main__":
# Inference(OnnxEngine('azaion-2025-03-10.onnx', batch_size=4),
# confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')
@@ -23,26 +32,19 @@ if __name__ == "__main__":
# Inference(TensorRTEngine('azaion-2025-03-10_batch8.engine', batch_size=8),
# confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')
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)
api_client = ApiClient()
key = Security.get_model_encryption_key()
engine_filename = TensorRTEngine.get_engine_filename()
model_bytes = api_client.load_big_small_resource(engine_filename, 'models', key)
cdn_manager = CDNManager(CDNCredentials(cdn_c.host,
cdn_c.downloader_access_key, cdn_c.downloader_access_secret,
cdn_c.uploader_access_key, cdn_c.uploader_access_secret))
cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG)
cdn_manager.download(cdn_c.bucket, constants.AI_ONNX_MODEL_FILE_BIG)
Inference(TensorRTEngine(model_bytes),
confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4')
api_client = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL)
onxx_model_bytes = api_client.load_resource(constants.AI_ONNX_MODEL_FILE_BIG, constants.AI_ONNX_MODEL_FILE_SMALL)
# cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG)
# tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL)
# Inference(OnnxEngine(onxx_model_bytes, batch_size=4),
# confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4')
Inference(TensorRTEngine(tensor_model_bytes, batch_size=4),
confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4')
train.py
+11 -8
View File
@@ -7,12 +7,13 @@ from datetime import datetime
from os import path, replace, listdir, makedirs, scandir
from os.path import abspath
from pathlib import Path
from time import sleep
import yaml
from ultralytics import YOLO
import constants
from azaion_api import ApiCredentials, Api
from api_client import ApiCredentials, ApiClient
from cdn_manager import CDNCredentials, CDNManager
from constants import (processed_images_dir,
processed_labels_dir,
@@ -20,10 +21,11 @@ from constants import (processed_images_dir,
datasets_dir, models_dir,
corrupted_images_dir, corrupted_labels_dir, sample_dir)
from dto.annotationClass import AnnotationClass
from inference.onnx_engine import OnnxEngine
from security import Security
from utils import Dotdict
from exports import export_tensorrt, upload_model
from exports import export_tensorrt, upload_model, export_onnx
today_folder = f'{prefix}{datetime.now():{date_format}}'
today_dataset = path.join(datasets_dir, today_folder)
@@ -224,10 +226,11 @@ if __name__ == '__main__':
# validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
# form_data_sample(500)
# convert2rknn()
model_path = 'azaion.pt'
export_tensorrt(model_path)
engine_model_path = f'{Path(model_path).stem}.engine'
upload_model(engine_model_path, engine_model_path)
api_client = ApiClient()
onnx_path = export_onnx('azaion.pt')
onnx_model_path = f'{Path(model_path).stem}.onnx'
upload_model(onnx_model_path, onnx_model_path)
with open(onnx_path, 'rb') as binary_file:
onnx_bytes = binary_file.read()
key = Security.get_model_encryption_key()
api_client.upload_big_small_resource(onnx_bytes, onnx_path, constants.MODELS_FOLDER, key)