确定trt模型的参数量

Question

将pytorch训练好的模型转换成trt格式后，怎么确定转换得到的trt文件包含的参数量呢？
还有怎么比较tensorrt的压缩优化效果呢？

Answer 1

引自免费微信小程序：皆我百晓生

要确定TRT（TensorRT）模型的参数量以及比较其压缩优化效果，可以按照以下步骤进行：

1. 确定TRT模型的参数量：

在PyTorch中，您可以使用torchsummary库来获取原始模型的参数量。对于已转换为TRT格式的模型，由于其内部结构不便于直接查询，我们可以通过以下间接方法估算参数量：

方法一：对比原始模型与TRT模型输出一致性

首先，在转换前保存原始PyTorch模型的权重，并确保TRT模型加载这些权重进行推理。然后，对两者进行相同的输入数据推断，并验证输出结果的一致性。若输出一致，则可认为TRT模型包含的参数量与原始模型相同。

以下是示例代码：

import torch
from torchsummary import summary

# 假设您已经训练好了一个名为model的PyTorch模型
# 保存原始模型权重
torch.save(model.state_dict(), 'original_model.pth')

# 将模型转换为TRT模型
# ...（此处省略具体的模型转换过程）

# 加载原始模型权重到TRT模型
trt_model.load_state_dict(torch.load('original_model.pth'))

# 检查原始模型与TRT模型输出一致性
input_data = ...  # 生成或准备测试用的输入数据
with torch.no_grad():
    orig_output = model(input_data)
    trt_output = trt_model(input_data)

    assert torch.allclose(orig_output, trt_output), "Outputs are not equal"

# 若通过一致性检查，可以认为TRT模型参数量与原始模型相同
print("Original PyTorch Model Parameters:")
summary(model, input_size=(...))  # 替换为您的模型输入尺寸

# 参数量即为summary输出中的"Total params"

方法二：查阅官方文档或社区资源

如果您使用的特定模型在TensorRT官方文档或社区中有明确记载其转换后的参数量信息，可以直接参考这些资料。

2. 比较TensorRT的压缩优化效果：

评估TensorRT的压缩优化效果通常从以下几个方面进行：

a) 模型大小比较：

测量原始PyTorch模型（.pth文件）和转换后的TRT模型（.engine文件）的文件大小。较小的文件大小通常意味着更高的压缩率和更优的存储效率。

ls -lh original_model.pth converted_trt_model.engine

b) 推理速度比较：

在相同的硬件环境下，分别运行原始PyTorch模型和TRT模型进行推理，并记录平均推理时间。更快的推理速度表示TensorRT优化效果显著。

import time

def benchmark(model, input_data):
    with torch.no_grad():
        start_time = time.time()
        for _ in range(num_runs):  # 设置合适的num_runs值，如100
            output = model(input_data)
        end_time = time.time()

    return (end_time - start_time) / num_runs

input_data = ...  # 生成或准备测试用的输入数据

orig_inference_time = benchmark(model, input_data)
trt_inference_time = benchmark(trt_model, input_data)

print(f"Original Model Inference Time: {orig_inference_time} seconds")
print(f"TRT Model Inference Time: {trt_inference_time} seconds")

c) 资源利用率比较：

监控并比较两种模型在推理过程中CPU、GPU内存使用情况及显存占用。更低的资源消耗表明TensorRT优化更高效。

可以使用系统监控工具（如nvidia-smi、top等）或编程接口（如torch.cuda.memory_allocated()、psutil库等）来收集相关数据。

总结来说，通过上述方法，您可以确定TRT模型的参数量，并从模型大小、推理速度、资源利用率等方面全面评估TensorRT的压缩优化效果。尽管无法直接查看TRT模型的参数量，但通过输出一致性检查和性能指标对比，可以有效判断其优化效果。