143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import time
|
|
import numpy as np
|
|
from resnet_attention import resnet18_cbam, resnet34_cbam, resnet50_cbam
|
|
|
|
# 设置随机种子以确保结果可复现
|
|
torch.manual_seed(42)
|
|
np.random.seed(42)
|
|
|
|
# 设备配置
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
print(f"测试设备: {device}")
|
|
|
|
# 测试参数
|
|
batch_sizes = [1, 4, 8, 16]
|
|
image_sizes = [224, 384, 512]
|
|
num_runs = 100 # 每个配置运行的次数
|
|
warmup_runs = 20 # 预热运行次数，排除启动开销
|
|
|
|
# 模型配置
|
|
model_configs = {
|
|
"resnet18": {
|
|
"base_model": lambda: resnet18_cbam(use_cbam=False),
|
|
"attention_model": lambda: resnet18_cbam(use_cbam=True)
|
|
},
|
|
"resnet34": {
|
|
"base_model": lambda: resnet34_cbam(use_cbam=False),
|
|
"attention_model": lambda: resnet34_cbam(use_cbam=True)
|
|
},
|
|
"resnet50": {
|
|
"base_model": lambda: resnet50_cbam(use_cbam=False),
|
|
"attention_model": lambda: resnet50_cbam(use_cbam=True)
|
|
}
|
|
}
|
|
|
|
|
|
# 基准测试函数
|
|
def benchmark_model(model, input_size, batch_size, num_runs, warmup_runs):
|
|
"""
|
|
测试模型的推理性能
|
|
|
|
参数:
|
|
- model: 待测试的模型
|
|
- input_size: 输入图像尺寸
|
|
- batch_size: 批次大小
|
|
- num_runs: 测试运行次数
|
|
- warmup_runs: 预热运行次数
|
|
|
|
返回:
|
|
- 平均推理时间(毫秒)
|
|
- 吞吐量(样本/秒)
|
|
"""
|
|
# 设置为评估模式
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
# 创建随机输入
|
|
input_tensor = torch.randn(batch_size, 3, input_size, input_size, device=device)
|
|
|
|
# 预热
|
|
with torch.no_grad():
|
|
for _ in range(warmup_runs):
|
|
_ = model(input_tensor)
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize() # 同步GPU操作
|
|
|
|
# 测量推理时间
|
|
start_time = time.time()
|
|
with torch.no_grad():
|
|
for _ in range(num_runs):
|
|
_ = model(input_tensor)
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize() # 同步GPU操作
|
|
end_time = time.time()
|
|
|
|
# 计算指标
|
|
total_time = end_time - start_time
|
|
avg_time_per_batch = total_time / num_runs * 1000 # 毫秒
|
|
throughput = batch_size * num_runs / total_time # 样本/秒
|
|
|
|
return avg_time_per_batch, throughput
|
|
|
|
|
|
# 运行测试
|
|
results = {}
|
|
|
|
for model_name, config in model_configs.items():
|
|
results[model_name] = {}
|
|
|
|
# 创建模型
|
|
base_model = config["base_model"]()
|
|
attention_model = config["attention_model"]()
|
|
|
|
# 计算参数量
|
|
base_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
|
|
attention_params = sum(p.numel() for p in attention_model.parameters() if p.requires_grad)
|
|
param_increase = (attention_params - base_params) / base_params * 100
|
|
|
|
print(f"\n测试模型: {model_name}")
|
|
print(f" 基础参数量: {base_params / 1e6:.2f}M")
|
|
print(f" 带注意力参数量: {attention_params / 1e6:.2f}M")
|
|
print(f" 参数量增加: {param_increase:.2f}%")
|
|
|
|
for batch_size in batch_sizes:
|
|
for image_size in image_sizes:
|
|
key = f"batch_{batch_size}_size_{image_size}"
|
|
results[model_name][key] = {}
|
|
|
|
# 测试基础模型
|
|
base_time, base_throughput = benchmark_model(
|
|
base_model, image_size, batch_size, num_runs, warmup_runs
|
|
)
|
|
|
|
# 测试注意力模型
|
|
attention_time, attention_throughput = benchmark_model(
|
|
attention_model, image_size, batch_size, num_runs, warmup_runs
|
|
)
|
|
|
|
# 计算增加的百分比
|
|
time_increase = (attention_time - base_time) / base_time * 100
|
|
throughput_decrease = (base_throughput - attention_throughput) / base_throughput * 100
|
|
|
|
results[model_name][key]["base_time"] = base_time
|
|
results[model_name][key]["attention_time"] = attention_time
|
|
results[model_name][key]["time_increase"] = time_increase
|
|
results[model_name][key]["base_throughput"] = base_throughput
|
|
results[model_name][key]["attention_throughput"] = attention_throughput
|
|
results[model_name][key]["throughput_decrease"] = throughput_decrease
|
|
|
|
print(f" 配置: 批次大小={batch_size}, 图像尺寸={image_size}x{image_size}")
|
|
print(f" 基础模型: 平均时间={base_time:.2f}ms, 吞吐量={base_throughput:.2f}样本/秒")
|
|
print(f" 注意力模型: 平均时间={attention_time:.2f}ms, 吞吐量={attention_throughput:.2f}样本/秒")
|
|
print(f" 时间增加: {time_increase:.2f}%, 吞吐量下降: {throughput_decrease:.2f}%")
|
|
|
|
# 保存结果
|
|
import json
|
|
|
|
with open('benchmark_results.json', 'w') as f:
|
|
json.dump(results, f, indent=2)
|
|
|
|
print("\n测试完成，结果已保存到 benchmark_results.json")
|