모델 경량화와 배포

import torch import onnx import onnxruntime as ort def export_to_onnx(model, save_path: str, input_size=(1, 3, 224, 224)): """PyTorch 모델을 ONNX로 변환""" model.eval() # 더미 입력 생성 dummy_input = torch.randn(input_size) # ONNX 내보내기 torch.onnx.export( model, dummy_input, save_path, export_params=True, opset_version=17, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, 'output': {0: 'batch_size'} } ) # 검증 onnx_model = onnx.load(save_path) onnx.checker.check_model(onnx_model) print(f"ONNX 모델 저장: {save_path}") class ONNXInference: """ONNX Runtime 추론 엔진""" def __init__(self, onnx_path: str, device='cuda'): providers = ['CUDAExecutionProvider'] if device == 'cuda' else ['CPUExecutionProvider'] self.session = ort.InferenceSession(onnx_path, providers=providers) self.input_name = self.session.get_inputs()[0].name def predict(self, image: np.ndarray) -> np.ndarray: """추론 실행""" outputs = self.session.run(None, {self.input_name: image}) return outputs[0] # 사용 예시 model = DefectClassifier(num_classes=5) model.load_state_dict(torch.load('best_model.pth')['model_state_dict']) export_to_onnx(model, 'model.onnx') # ONNX 추론 테스트 engine = ONNXInference('model.onnx', device='cuda') result = engine.predict(test_image)

import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit class TensorRTEngine: """TensorRT 추론 엔진""" def __init__(self, engine_path: str): self.logger = trt.Logger(trt.Logger.WARNING) with open(engine_path, 'rb') as f: self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() # 메모리 할당 self.inputs, self.outputs, self.bindings, self.stream = self._allocate_buffers() def _allocate_buffers(self): inputs, outputs, bindings = [], [], [] stream = cuda.Stream() for binding in self.engine: size = trt.volume(self.engine.get_binding_shape(binding)) dtype = trt.nptype(self.engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) bindings.append(int(device_mem)) if self.engine.binding_is_input(binding): inputs.append({'host': host_mem, 'device': device_mem}) else: outputs.append({'host': host_mem, 'device': device_mem}) return inputs, outputs, bindings, stream def infer(self, image: np.ndarray) -> np.ndarray: """TensorRT 추론""" # 입력 복사 np.copyto(self.inputs[0]['host'], image.ravel()) cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream) # 추론 실행 self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle) # 출력 복사 cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream) self.stream.synchronize() return self.outputs[0]['host'].copy() # ONNX → TensorRT 변환 def build_tensorrt_engine(onnx_path: str, engine_path: str, precision='fp16'): """ONNX를 TensorRT 엔진으로 변환""" logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) # ONNX 파싱 with open(onnx_path, 'rb') as f: if not parser.parse(f.read()): for error in range(parser.num_errors): print(parser.get_error(error)) raise RuntimeError("ONNX 파싱 실패") # 빌더 설정 config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB if precision == 'fp16': config.set_flag(trt.BuilderFlag.FP16) # 엔진 빌드 engine = builder.build_serialized_network(network, config) with open(engine_path, 'wb') as f: f.write(engine) print(f"TensorRT 엔진 저장: {engine_path}")

# PyTorch 동적 양자화 import torch.quantization as quant def quantize_model_dynamic(model): """동적 양자화 (FC 레이어 대상)""" model.eval() quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 ) return quantized_model # TensorRT INT8 캘리브레이션 class ImageCalibrator(trt.IInt8EntropyCalibrator2): """INT8 양자화용 캘리브레이션 데이터 제공""" def __init__(self, dataloader, cache_file='calibration.cache'): super().__init__() self.dataloader = iter(dataloader) self.cache_file = cache_file self.batch_size = dataloader.batch_size self.current_batch = 0 # GPU 메모리 할당 batch = next(iter(dataloader))[0] self.device_input = cuda.mem_alloc(batch.numpy().nbytes) def get_batch_size(self): return self.batch_size def get_batch(self, names): try: batch = next(self.dataloader)[0] cuda.memcpy_htod(self.device_input, batch.numpy()) return [int(self.device_input)] except StopIteration: return None def read_calibration_cache(self): if os.path.exists(self.cache_file): with open(self.cache_file, 'rb') as f: return f.read() return None def write_calibration_cache(self, cache): with open(self.cache_file, 'wb') as f: f.write(cache)

import time import logging class ProductionInferenceServer: """생산 환경 추론 서버""" def __init__(self, model_path: str, config: dict): self.config = config self.logger = logging.getLogger(__name__) # 모델 로드 self.engine = TensorRTEngine(model_path) # 전처리 파이프라인 self.preprocessor = ImagePreprocessor( target_size=config['input_size'] ) # 성능 모니터링 self.inference_times = [] self.total_inferences = 0 def predict(self, image: np.ndarray) -> dict: """추론 실행""" start_time = time.time() # 전처리 processed = self.preprocessor.preprocess(image) # 추론 output = self.engine.infer(processed) probs = self._softmax(output) # 결과 해석 pred_class = int(np.argmax(probs)) confidence = float(probs[pred_class]) # 성능 기록 inference_time = (time.time() - start_time) * 1000 self.inference_times.append(inference_time) self.total_inferences += 1 return { 'class': self.config['class_names'][pred_class], 'class_id': pred_class, 'confidence': confidence, 'is_defect': pred_class != 0, 'inference_time_ms': inference_time, 'all_probs': probs.tolist() } def get_metrics(self) -> dict: """성능 메트릭 반환""" if not self.inference_times: return {} times = np.array(self.inference_times[-1000:]) # 최근 1000건 return { 'total_inferences': self.total_inferences, 'avg_inference_ms': float(np.mean(times)), 'p95_inference_ms': float(np.percentile(times, 95)), 'p99_inference_ms': float(np.percentile(times, 99)), 'fps': 1000 / np.mean(times) } @staticmethod def _softmax(x): e_x = np.exp(x - np.max(x)) return e_x / e_x.sum()