更新性能测试脚本

tests/performance_loss_test.py

@@ -52,7 +52,13 @@ class UsageMonitor:
         return avg_cpu_usage, avg_gpu_usage
 
 
-def process_image(image_path):
+def process_image(image_path, transpose=True):
+    """
+    图片处理
+    :param image_path: 图片路径
+    :param transpose: 是否进行维度转换，在使用pytorch框架训练出来的权重需要进行维度转换，tensorflow、keras框架不需要
+    :return:
+    """
     # 打开图像并转换为RGB
     image = Image.open(image_path).convert("RGB")
 
@@ -66,18 +72,20 @@ def process_image(image_path):
     mean = np.array([0.485, 0.456, 0.406])
     std = np.array([0.229, 0.224, 0.225])
     image_array = (image_array - mean) / std
-    image_array = image_array.transpose((2, 0, 1)).copy()
+    if transpose:
+        image_array = image_array.transpose((2, 0, 1)).copy()
 
     return image_array.astype(np.float32)
 
 
-def batch_predict_images(session, image_dir, target_class, batch_size=10):
+def batch_predict_images(session, image_dir, target_class, batch_size=10, pytorch=True):
     """
     对指定图片文件夹图片进行批量检测
     :param session: onnx runtime session
     :param image_dir: 待推理的图像文件夹
     :param target_class: 目标分类
-    :param batch_size: 每批图片数量
+    :param batch_size: 每批图片数量, 默认为10
+    :param pytorch: 模型是否使用pytorch框架训练出的权重导出的onnx文件，默认为True
     :return: 检测结果
     """
     image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
@@ -92,7 +100,7 @@ def batch_predict_images(session, image_dir, target_class, batch_size=10):
 
         for image_file in batch_files:
             image_path = os.path.join(image_dir, image_file)
-            image = process_image(image_path)
+            image = process_image(image_path, pytorch)
             batch_images.append(image)
 
         # 将批次图片堆叠成 (batch_size, 3, 224, 224) 维度
@@ -119,15 +127,17 @@ def batch_predict_images(session, image_dir, target_class, batch_size=10):
 # 模型推理函数
 def model_inference(model_filename, val_dataset_dir):
     """
-    模型推理
+    模型推理验证集目录下所有图片
     :param model_filename: 模型文件
     :param val_dataset_dir: 验证集图片目录
     :return: 验证集推理准确率
     """
-    if ort.get_available_providers():
-        session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])
-    else:
-        session = ort.InferenceSession(model_filename)
+    # 以下使用GPU进行推理出现问题，需要较新的CUDA版本，默认使用CPU进行推理
+    # if ort.get_available_providers():
+    #     session = ort.InferenceSession(model_filename, providers=['CUDAExecutionProvider'])
+    # else:
+    #     session = ort.InferenceSession(model_filename)
+    session = ort.InferenceSession(model_filename)
     accuracy = 0
     class_num = 0
     index = 0
@@ -137,7 +147,8 @@ def model_inference(model_filename, val_dataset_dir):
         if not os.path.isdir(class_path):
             continue
         class_num += 1
-        batch_result = batch_predict_images(session, class_path, index)
+        is_pytorch = False if "keras" in model_filename or "tensorflow" in model_filename else True
+        batch_result = batch_predict_images(session, class_path, index, pytorch=is_pytorch)
         accuracy += batch_result
         index += 1
     print(f"class_num: {class_num}, index: {index}")
@@ -146,12 +157,17 @@ def model_inference(model_filename, val_dataset_dir):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='模型推理性能验证脚本')
-    # parser.add_argument('--model_file', default=None, type=str, help='待测试模型的onnx文件')
+    # parser.add_argument('--origin_model_file', default=None, type=str, help='待测试原始模型的onnx文件')
+    # parser.add_argument('--watermark_model_file', default=None, type=str, help='待测试水印模型的onnx文件')
     # parser.add_argument('--val_dataset_dir', default=None, type=str, help='验证集目录')
-    parser.add_argument('--model_file', default="origin_models/googlenet/model_139.onnx", type=str, help='待测试模型的onnx文件')
+
+    parser.add_argument('--origin_model_file', default="origin_models/alexnet_keras/alexnet_090.onnx", type=str,
+                        help='待测试原始模型的onnx文件')
+    parser.add_argument('--watermark_model_file', default="test_models/alexnet_keras/alexnet_090.onnx", type=str,
+                        help='待测试水印模型的onnx文件')
     parser.add_argument('--val_dataset_dir', default="val", type=str, help='验证集目录')
     args, _ = parser.parse_known_args()
-    if args.model_file is None:
+    if args.origin_model_file is None:
         raise Exception("待测试模型的onnx文件不可为空")
     if args.val_dataset_dir is None:
         raise Exception("验证集目录不可为空")
@@ -161,14 +177,38 @@ if __name__ == '__main__':
     # 记录推理开始时间
     start_time = time.time()
     # 进行模型推理
-    accuracy = model_inference(args.model_file, args.val_dataset_dir)
+    accuracy = model_inference(args.origin_model_file, args.val_dataset_dir)
     # 记录推理结束时间
     end_time = time.time()
     monitor.stop()
     # 输出平均 CPU 和 GPU 使用率
     avg_cpu, avg_gpu = monitor.get_average_usage()
-
+    print("原始模型推理性能：")
     print(f"平均 CPU 使用率：{avg_cpu:.2f}%")
     print(f"平均 GPU 使用率：{avg_gpu:.2f}%")
     print(f"模型推理时间: {end_time - start_time:.2f} 秒")
     print(f"准确率: {accuracy * 100:.2f}%")
+
+    if args.watermark_model_file:  # 加入存在比对模型，进行再次推理，然后统计性能指标
+        time.sleep(20)
+        monitor2 = UsageMonitor(interval=0.5)  # 每隔 0.5 秒采样一次
+        monitor2.start()
+        # 记录推理开始时间
+        start_time2 = time.time()
+        # 进行模型推理
+        accuracy2 = model_inference(args.watermark_model_file, args.val_dataset_dir)
+        # 记录推理结束时间
+        end_time2 = time.time()
+        monitor2.stop()
+        # 输出平均 CPU 和 GPU 使用率
+        avg_cpu2, avg_gpu2 = monitor2.get_average_usage()
+        print("水印模型推理性能：")
+        print(f"平均 CPU 使用率：{avg_cpu2:.2f}%")
+        print(f"平均 GPU 使用率：{avg_gpu2:.2f}%")
+        print(f"模型推理时间: {end_time2 - start_time2:.2f} 秒")
+        print(f"准确率: {accuracy2 * 100:.2f}%")
+
+        print("------------------性能指标如下-------------------------")
+        print(f"嵌入后模型推理准确率下降值：{(accuracy - accuracy2) * 100:.2f}%")
+        print(f"算力资源消耗增加值：{(avg_cpu2 - avg_cpu):.2f}%")
+        print(f"运行效率降低值: {((end_time2 - start_time2) - (end_time - start_time)) * 100 / (end_time - start_time):.2f} %")

tests/prune_tool.py

@@ -63,7 +63,7 @@ def find_onnx_files(root_dir):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='模型文件剪枝工具')
     # parser.add_argument('--target_dir', default=None, type=str, help='待剪枝的模型文件存放根目录，支持子文件夹递归处理')
-    parser.add_argument('--target_dir', default="origin_models/googlenet", type=str, help='待剪枝的模型文件存放根目录，支持子文件夹递归处理')
+    parser.add_argument('--target_dir', default="blackbox_models/alexnet_keras", type=str, help='待剪枝的模型文件存放根目录，支持子文件夹递归处理')
     parser.add_argument('--pruned_saved_dir', default=None, type=str, help='剪枝模型文件保存目录，默认为None，表示与原始onnx权重文件放在同一目录下')
     parser.add_argument('--percent', default=0.05, type=float, help='规则剪枝百分比')
     args, _ = parser.parse_known_args()

tests/verify_tool_accuracy_test.py

@@ -5,7 +5,7 @@ from watermark_verify import verify_tool
 
 model_types = {
     "classification": [
-        "alexnet", "googlenet", "resnet", "vgg16"
+        "alexnet","alexnet_keras", "googlenet", "resnet", "vgg16", "vgg16_tensorflow"
     ],
     "object_detection": [
         "ssd", "yolox", "rcnn"
@@ -38,10 +38,10 @@ if __name__ == '__main__':
     # parser.add_argument('--except_result', default=None, type=str, help='模型推理预期结果。默认为None')
 
     parser.add_argument('--model_type', default="classification", type=str, help='按照模型分类过滤，用于区分是目标检测模型还是图像分类模型，可选参数：classification、objection_detect')
-    parser.add_argument('--model_value', default="googlenet", type=str, help='按照模型名称过滤，可选参数：alexnet、googlenet、resnet、vgg16、ssd、yolox、rcnn')
-    parser.add_argument('--target_dir', default="origin_models", type=str,
+    parser.add_argument('--model_value', default="alexnet_keras", type=str, help='按照模型名称过滤，可选参数：alexnet、googlenet、resnet、vgg16、ssd、yolox、rcnn')
+    parser.add_argument('--target_dir', default="blackbox_models", type=str,
                         help='模型文件存放根目录，支持子文件夹递归处理')
-    parser.add_argument('--except_result', default="False", type=str, help='模型推理预期结果。默认为None')
+    parser.add_argument('--except_result', default="True", type=str, help='模型推理预期结果。默认为None')
     parser.add_argument('--model_file_filter', default="pruned", type=str,
                         help='按照模型文件名过滤, 比如剪枝模型文件名存在pruned。默认为None')
 
@@ -69,6 +69,8 @@ if __name__ == '__main__':
         onnx_files = [os.path.abspath(item) for item in onnx_files]
         if args.model_file_filter:
             onnx_files = [item for item in onnx_files if args.model_file_filter in item]
+        else:
+            onnx_files = [item for item in onnx_files if "pruned" not in item]
         print(f"model_name: {model_dir}\nonnx_files:")
         print(*onnx_files, sep='\n')
         for onnx_file in onnx_files: