教程 7：实用工具（待更新）¶

我们在 tools/ 目录下提供了很多有用的工具。

获取 FLOP 和参数量（实验性）¶

我们提供了一个改编自 flops-counter.pytorch 的脚本来计算模型的 FLOP 和参数量。

python tools/get_flops.py ${CONFIG_FILE} [--shape ${INPUT_SHAPE}]

例如，

python tools/get_flops.py configs/resotorer/srresnet.py --shape 40 40

你会得到以下的结果。

==============================
Input shape: (3, 40, 40)
Flops: 4.07 GMac
Params: 1.52 M
==============================

注：此工具仍处于实验阶段，我们不保证数字正确。您可以将结果用于简单的比较，但在技术报告或论文中采用它之前，请仔细检查它。

(1) FLOPs 与输入形状有关，而参数量与输入形状无关。默认输入形状为 (1, 3, 250, 250)。 (2) 一些运算符不计入 FLOP，如 GN 和自定义运算符。你可以通过修改 mmcv/cnn/utils/flops_counter.py 来添加对新运算符的支持。

发布模型¶

在将模型上传到 AWS 之前，您可能需要 (1) 将模型权重转换为 CPU tensors, (2) 删除优化器状态，和 (3) 计算模型权重文件的哈希并将哈希 ID 附加到文件名。

python tools/publish_model.py ${INPUT_FILENAME} ${OUTPUT_FILENAME}

例如，

python tools/publish_model.py work_dirs/example_exp/latest.pth example_model_20200202.pth

最终输出文件名将是 example_model_20200202-{hash id}.pth.

转换为 ONNX（实验性）¶

我们提供了一个脚本将模型转换为 ONNX 格式。转换后的模型可以通过 Netron 等工具进行可视化。此外，我们还支持比较 Pytorch 和 ONNX 模型之间的输出结果。

python tools/pytorch2onnx.py
    ${CFG_PATH} \
    ${CHECKPOINT_PATH} \
    ${MODEL_TYPE} \
    ${IMAGE_PATH} \
    --trimap-path ${TRIMAP_PATH} \
    --output-file ${OUTPUT_ONNX} \
    --show \
    --verify \
    --dynamic-export

参数说明：

config : 模型配置文件的路径。
checkpoint : 模型模型权重文件的路径。
model_type : 配置文件的模型类型，选项： inpainting, mattor, restorer, synthesizer。
image_path : 输入图像文件的路径。
--trimap-path : 输入三元图文件的路径，用于 mattor 模型。
--output-file: 输出 ONNX 模型的路径。默认为 tmp.onnx。
--opset-version : ONNX opset 版本。默认为 11。
--show: 确定是否打印导出模型的架构。默认为 False。
--verify: 确定是否验证导出模型的正确性。默认为 False。
--dynamic-export: 确定是否导出具有动态输入和输出形状的 ONNX 模型。默认为 False。

注：此工具仍处于试验阶段。目前不支持某些自定义运算符。我们现在只支持 mattor 和 restorer。

支持导出到 ONNX 的模型列表¶

下表列出了保证可导出到 ONNX 并可在 ONNX Runtime 中运行的模型。

模型	配置	动态形状	批量推理
ESRGAN	esrgan_x4c64b23g32_g1_400k_div2k.py	Y	Y
ESRGAN	esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py	Y	Y
SRCNN	srcnn_x4k915_g1_1000k_div2k.py	Y	Y
DIM	dim_stage3_v16_pln_1x1_1000k_comp1k.py	Y	Y
GCA	gca_r34_4x10_200k_comp1k.py	N	Y
IndexNet	indexnet_mobv2_1x16_78k_comp1k.py	Y	Y

注：

以上所有模型均使用 Pytorch==1.6.0 和 onnxruntime==1.5.1
如果您遇到上面列出的模型的任何问题，请创建一个 issue，我们会尽快处理。对于列表中未包含的型号，请尝试自行解决。
由于此功能是实验性的并且可能会快速更改，请始终尝试使用最新的 mmcv 和 mmagic。

将 ONNX 转换为 TensorRT（实验性）¶

我们还提供了将 ONNX 模型转换为 TensorRT 格式的脚本。此外，我们支持比较 ONNX 和 TensorRT 模型之间的输出结果。

python tools/onnx2tensorrt.py
    ${CFG_PATH} \
    ${MODEL_TYPE} \
    ${IMAGE_PATH} \
    ${INPUT_ONNX} \
    --trt-file ${OUT_TENSORRT} \
    --max-shape INT INT INT INT \
    --min-shape INT INT INT INT \
    --workspace-size INT \
    --fp16 \
    --show \
    --verify \
    --verbose

参数说明：

config : 模型配置文件的路径。
model_type :配置文件的模型类型，选项： inpainting, mattor, restorer, synthesizer。
img_path : 输入图像文件的路径。
onnx_file : 输入 ONNX 文件的路径。
--trt-file : 输出 TensorRT 模型的路径。默认为 tmp.trt。
--max-shape : 模型输入的最大形状。
--min-shape : 模型输入的最小形状。
--workspace-size: 以 GiB 为单位的最大工作空间大小。默认为 1 GiB。
--fp16: 确定是否以 fp16 模式导出 TensorRT。默认为 False。
--show: 确定是否显示 ONNX 和 TensorRT 的输出。默认为 False。
--verify: 确定是否验证导出模型的正确性。默认为 False。
--verbose: 确定在创建 TensorRT 引擎时是否详细记录日志消息。默认为 False。

注：此工具仍处于试验阶段。目前不支持某些自定义运算符。我们现在只支持 restorer。在生成 SRCNN 的 ONNX 文件时，将 SCRNN 模型中的 ‘bicubic’ 替换为 ‘bilinear’ [此处](https://github.com/open-mmlab/mmagic/blob/764e6065e315b7d0033762038fcbf0bb1c570d4d/mmagic.bones/modelsrnn py#L40）。因为 TensorRT 目前不支持 bicubic 插值，最终性能将下降约 4%。

支持导出到 TensorRT 的模型列表¶

下表列出了保证可导出到 TensorRT 引擎并可在 TensorRT 中运行的模型。

模型	配置	动态形状	批量推理	备注
ESRGAN	esrgan_x4c64b23g32_g1_400k_div2k.py	Y	Y
ESRGAN	esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py	Y	Y
SRCNN	srcnn_x4k915_g1_1000k_div2k.py	Y	Y	'bicubic' 上采样必须替换为 'bilinear'

注：

以上所有模型均使用 Pytorch==1.8.1、onnxruntime==1.7.0 和 tensorrt==7.2.3.4 进行测试
如果您遇到上面列出的模型的任何问题，请创建一个问题，我们会尽快处理。对于列表中未包含的型号，请尝试自行解决。
由于此功能是实验性的并且可能会快速更改，因此请始终尝试使用最新的 mmcv 和 mmagic。

评估 ONNX 和 TensorRT 模型（实验性）¶

我们在 tools/deploy_test.py 中提供了评估 TensorRT 和 ONNX 模型的方法。

先决条件¶

要评估 ONNX 和 TensorRT 模型，应先安装 onnx、onnxruntime 和 TensorRT。遵循 mmcv 中的 ONNXRuntime 和 [mmcv 中的 TensorRT 插件](https://github.com/open-mmlab/mmcv/blob/master/docs/tensorrt_plugin.md%EF%BC%89%E4%BD%BF%E7%94%A8 ONNXRuntime 自定义操作和 TensorRT 插件安装 mmcv-full。

用法¶

python tools/deploy_test.py \
    ${CONFIG_FILE} \
    ${MODEL_PATH} \
    ${BACKEND} \
    --out ${OUTPUT_FILE} \
    --save-path ${SAVE_PATH} \
    ----cfg-options ${CFG_OPTIONS} \

参数说明：¶

config: 模型配置文件的路径。
model: TensorRT 或 ONNX 模型文件的路径。
backend: 用于测试的后端，选择 tensorrt 或 onnxruntime。
--out: pickle 格式的输出结果文件的路径。
--save-path: 存储图像的路径，如果没有给出，则不会保存图像。
--cfg-options: 覆盖使用的配置文件中的一些设置，xxx=yyy 格式的键值对将被合并到配置文件中。

结果和模型¶

Model	Config	Dataset	Metric	PyTorch	ONNX Runtime	TensorRT FP32	TensorRT FP16
ESRGAN	`esrgan_x4c64b23g32_g1_400k_div2k.py`	Set5	PSNR	28.2700	28.2619	28.2619	28.2616
		Set5	SSIM	0.7778	0.7784	0.7784	0.7783
		Set14	PSNR	24.6328	24.6290	24.6290	24.6274
		Set14	SSIM	0.6491	0.6494	0.6494	0.6494
		DIV2K	PSNR	26.6531	26.6532	26.6532	26.6532
		DIV2K	SSIM	0.7340	0.7340	0.7340	0.7340
ESRGAN	`esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py`	Set5	PSNR	30.6428	30.6307	30.6307	30.6305
		Set5	SSIM	0.8559	0.8565	0.8565	0.8566
		Set14	PSNR	27.0543	27.0422	27.0422	27.0411
		Set14	SSIM	0.7447	0.7450	0.7450	0.7449
		DIV2K	PSNR	29.3354	29.3354	29.3354	29.3339
		DIV2K	SSIM	0.8263	0.8263	0.8263	0.8263
SRCNN	`srcnn_x4k915_g1_1000k_div2k.py`	Set5	PSNR	28.4316	28.4120	27.2144	27.2127
		Set5	SSIM	0.8099	0.8106	0.7782	0.7781
		Set14	PSNR	25.6486	25.6367	24.8613	24.8599
		Set14	SSIM	0.7014	0.7015	0.6674	0.6673
		DIV2K	PSNR	27.7460	27.7460	26.9891	26.9862
		DIV2K	SSIM	0.7854	0.78543	0.7605	0.7604

注：

所有 ONNX 和 TensorRT 模型都使用数据集上的动态形状进行评估，图像根据原始配置文件进行预处理。
此工具仍处于试验阶段，我们目前仅支持 restorer。