azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 11:16: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
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')