Onnx 모델과 pth 모델 성능 비교

최근에 업무를 진행하여 npu 를 위한 모델 변환을 진행함

모델 변환 후 평가를 진행해보니, 기존 서버에서 측정한 모델 성능이 재현되지 않음

문제가 발생했을 가능성이 있는 부분은 2곳

1. 변환된 모델에 입력으로 넣어준 데이터
2. 변환하기 전 모델

일단 변환하기 전 onnx 모델을 살펴보기 위해 onnxruntime 으로 평가를 진행해보니 이 또한 서버 성능을 재현하지 못함.

import numpy as np
import onnxruntime as ort
import os
import glob

bin_directory = [데이터 디렉토리]

model_path = [모델 경로]

session = ort.InferenceSession(model_path)
input_name = session.get_inputs()[0].name

class_names = [클래스명 나열]

bin_files = glob.glob(os.path.join(bin_directory, "*.bin"))

results_list = []

for bin_file in bin_files:
    try:
        file_name = os.path.basename(bin_file)
        
        input_data = np.fromfile(bin_file, dtype=np.float32)
        
        input_data = input_data.reshape([입력 shape])
        
        results = session.run(None, {input_name: input_data})
        prediction = results[0]
        
        predicted_class_idx = np.argmax(prediction)
        predicted_class_prob = prediction[0][predicted_class_idx]
        predicted_class_name = class_names[predicted_class_idx] if predicted_class_idx < len(class_names) else f"미지정 클래스 {predicted_class_idx}"
        
        print(f"파일: {file_name}")
        print(f"예측 결과: {predicted_class_name} (인덱스: {predicted_class_idx}, 확률: {predicted_class_prob:.4f})")
        print(f"전체 확률 분포: {prediction[0]}")
        print("-" * 50)
        
        results_list.append({
            "file_name": file_name,
            "predicted_class": predicted_class_idx,
            "predicted_class_name": predicted_class_name,
            "probability": predicted_class_prob,
            "all_probabilities": prediction[0].tolist()
        })
        
    except Exception as e:
        print(f"오류 발생 - 파일: {bin_file}")
        print(f"오류 메시지: {str(e)}")
        print("-" * 50)

print(f"처리된 파일 수: {len(results_list)} / {len(bin_files)}")

 

따라서 다시 pth 를 onnx 로 변환하는 곳으로 돌아가서 문제가 있는지 확인

이 과정에서 랜덤한 텐서를 생성하여 pth 모델과 onnx 모델에 입력으로 주고 output 차이가 어떤지 확인

import torch
import torch.nn as nn
import numpy as np

def convert_to_onnx(
    model_path: str,
    config_path: str,
    save_path: str,
    input_shape: tuple = [입력 shape],
    seed: int = 42 
):
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    cfg = Config.fromfile(config_path)
    
    model = build_model(cfg.model)
    
    load_checkpoint(model, model_path, map_location='cpu')
    
    model.eval()
    
    class ModelWrapper(torch.nn.Module):
        def __init__(self, model):
            super(ModelWrapper, self).__init__()
            self.model = model
            self.softmax = nn.Softmax(dim=1) 
        def forward(self, x):
            with torch.no_grad():
                feat = self.model.extract_feat(x)  
                cls_score = self.model.cls_head(feat)
                return self.softmax(cls_score) 
    
    wrapped_model = ModelWrapper(model)
    
    dummy_input = torch.randn(input_shape)
    
    torch.onnx.export(
        wrapped_model,
        dummy_input,
        save_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=None,
        opset_version=11,
        do_constant_folding=True,
        verbose=True
    )
    
    print(f"Model has been converted to ONNX and saved at: {save_path}")
    
    import onnx
    onnx_model = onnx.load(save_path)
    onnx.checker.check_model(onnx_model, True)
    print("ONNX model checked successfully with full validation!")
    
    try:
        import onnxruntime as ort
        
        session = ort.InferenceSession(save_path)
        input_name = session.get_inputs()[0].name
        
        test_input = np.random.randn(*input_shape).astype(np.float32)
        
        onnx_outputs = session.run(None, {input_name: test_input})
        
        torch_input = torch.tensor(test_input)
        with torch.no_grad():
            feat = model.extract_feat(torch_input)
            torch_logits = model.cls_head(feat).numpy()
        
        with torch.no_grad():
            wrapped_output = wrapped_model(torch_input).numpy()
        
        print("\n===== 출력 비교 =====")
            
        print(f"\n2. PyTorch 모델 (probs):")
        for i, prob in enumerate(wrapped_output[0]):
            print(f"   클래스 {i}: {prob:.6f}")
            
        print(f"\n3. ONNX 모델 (probs):")
        for i, prob in enumerate(onnx_outputs[0][0]):
            print(f"   클래스 {i}: {prob:.6f}")
        
        print("\n===== 차이 분석 =====")
        max_diff_wrapped_onnx = np.max(np.abs(wrapped_output - onnx_outputs[0]))
        print(f"래핑된 PyTorch vs ONNX 최대 차이: {max_diff_wrapped_onnx:.6f}")
        
        # 합계 확인 (모두 약 1.0이어야 함)
        print("\n===== 확률 합계 확인 =====")
        print(f"래핑된 PyTorch 합계: {wrapped_output.sum():.6f}")
        print(f"ONNX 모델 합계: {onnx_outputs[0].sum():.6f}")
        
    except ImportError:
        print("onnxruntime is not installed. Skipping test inference.")

if __name__ == "__main__":
    model_path = [pth 모델 경로]
    config_path = [config path]
    save_path = [onnx 모델 저장 경로]
    
    # 모델 변환 실행 (시드 값 지정 가능)
    convert_to_onnx(
        model_path=model_path,
        config_path=config_path,
        save_path=save_path,
        seed=42  # 원하는 시드 값 지정
    )

---

결론: npu 용으로 모델 변환 전 onnx 로 변환이 잘 되었는지 확인하자