nnUNetV2转onnx格式

nnUNetV2转onnx，注意是V2格式，V1不支持！！！

假如已经训练好了nnUNetV2模型，使用五折交叉验证的方式完成！

训练过程可以参考：一文讲清nnUNetv2在自己数据集上训练

1. 转换代码：

import json
from os.path import isdir, join
from pathlib import Path
from typing import Tuple, Union

import numpy as np
import onnx
import onnxruntime
import torch
import shutil
from batchgenerators.utilities.file_and_folder_operations import load_json

from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
from nnunetv2.utilities.dataset_name_id_conversion import \
    maybe_convert_to_dataset_name
from nnunetv2.utilities.file_path_utilities import get_output_folder


def export_onnx_model(
    dataset_name_or_id: Union[int, str],
    output_dir: Path,
    configurations: Tuple[str] = (
        "2d",
        "3d_lowres",
        "3d_fullres",
        "3d_cascade_fullres",
    ),
    batch_size: int = 0,
    trainer: str = "nnUNetTrainer",
    plans_identifier: str = "nnUNetPlans",
    folds: Tuple[Union[int, str], ...] = (0, 1, 2, 3, 4),

    strict: bool = True,
    save_checkpoints: Tuple[str, ...] = ("checkpoint_final.pth",),
    output_names: tuple[str, ...] = None,
    verbose: bool = False,
) -> None:
    if not output_names:
        output_names = (f"{checkpoint[:-4]}.onnx" for checkpoint in save_checkpoints)

    if batch_size < 0:
        raise ValueError("batch_size must be non-negative")

    use_dynamic_axes = batch_size == 0

    dataset_name = maybe_convert_to_dataset_name(dataset_name_or_id)
    for c in configurations:
        print(f"Configuration {c}")
        trainer_output_dir = get_output_folder(
            dataset_name, trainer, plans_identifier, c
        )
        dataset_json = load_json(join(trainer_output_dir, "dataset.json"))

        # While we load in this file indirectly, we need the plans file to
        # determine the foreground intensity properties.
        plans = load_json(join(trainer_output_dir, 'plans.json'))
        foreground_intensity_properties = plans['foreground_intensity_properties_per_channel']

        if not isdir(trainer_output_dir):
            if strict:
                raise RuntimeError(
                    f"{dataset_name} is missing the trained model of configuration {c}"
                )
            else:
                print(f"Skipping configuration {c}, does not exist")
                continue

        predictor = nnUNetPredictor(
            # perform_everything_on_gpu=False,
            # device=torch.device("cpu"),
        )

        for checkpoint_name, output_name in zip(save_checkpoints, output_names):
            predictor.initialize_from_trained_model_folder(
                model_training_output_dir=trainer_output_dir,
                use_folds=folds,
                checkpoint_name=checkpoint_name,
                # disable_compilation=True,
            )

            list_of_parameters = predictor.list_of_parameters
            network = predictor.network
            config = predictor.configuration_manager

            for fold, params in zip(folds, list_of_parameters):
                network.load_state_dict(params)

                network.eval()

                # curr_output_dir = output_dir / c / f"fold_{fold}"
                curr_output_dir = Path(output_dir)
                if curr_output_dir.exists():
                    shutil.rmtree(curr_output_dir)  # 如果目录存在，删除目录及其所有内容
                curr_output_dir.mkdir(parents=True, exist_ok=True)
                # curr_output_dir = Path(output_dir)
                # if not curr_output_dir.exists():
                #     curr_output_dir.mkdir(parents=True)
                # else:
                #     if len(list(curr_output_dir.iterdir())) > 0:
                #         raise RuntimeError(
                #             f"Output directory {curr_output_dir} is not empty"
                #         )

                if use_dynamic_axes:
                    rand_input = torch.rand((1, 1, *config.patch_size))
                    torch_output = network(rand_input)

                    torch.onnx.export(
                        network,
                        rand_input,
                        curr_output_dir / output_name,
                        export_params=True,
                        verbose=verbose,
                        input_names=["input"],
                        output_names=["output"],
                        dynamic_axes={
                            "input": {0: "batch_size"},
                            "output": {0: "batch_size"},
                        },
                    )
                else:
                    rand_input = torch.rand((batch_size, 3, *config.patch_size))
                    torch_output = network(rand_input)

                    torch.onnx.export(
                        network,
                        rand_input,
                        curr_output_dir / output_name,
                        export_params=True,
                        verbose=verbose,
                        input_names=["input"],
                        output_names=["output"],
                    )

                onnx_model = onnx.load(curr_output_dir / output_name)
                onnx.checker.check_model(onnx_model)

                ort_session = onnxruntime.InferenceSession(
                    curr_output_dir / output_name, providers=["CPUExecutionProvider"]
                )
                ort_inputs = {ort_session.get_inputs()[0].name: rand_input.numpy()}
                ort_outs = ort_session.run(None, ort_inputs)

                try:
                    np.testing.assert_allclose(
                        torch_output.detach().cpu().numpy(),
                        ort_outs[0],
                        rtol=1e-03,
                        atol=1e-05,
                        verbose=True,
                    )
                except AssertionError as e:
                    print("WARN: Differences found between torch and onnx:\n")
                    print(e)
                    print(
                        "\nExport will continue, but please verify that your pipeline matches the original."
                    )

                print(f"Exported {curr_output_dir / output_name}")

                with open(curr_output_dir / "config.json", "w") as f:
                    config_dict = {
                        "configuration": c,
                        "fold": fold,
                        "model_parameters": {
                            "batch_size": batch_size
                            if not use_dynamic_axes
                            else "dynamic",
                            "patch_size": config.patch_size,
                            "spacing": config.spacing,
                            "normalization_schemes": config.normalization_schemes,
                            # These are mostly interesting for certification
                            # uses, but they are also useful for debugging.
                            # "UNet_class_name": config.UNet_class_name,
                            # "UNet_base_num_features": config.UNet_base_num_features,
                            # "unet_max_num_features": config.unet_max_num_features,
                            # "conv_kernel_sizes": config.conv_kernel_sizes,
                            # "pool_op_kernel_sizes": config.pool_op_kernel_sizes,
                            # "num_pool_per_axis": config.num_pool_per_axis,
                        },
                        "dataset_parameters": {
                            "dataset_name": dataset_name,
                            "num_channels": len(dataset_json["channel_names"].keys()),
                            "channels": {
                                k: {
                                    "name": v,
                                    # For when normalization is not Z-Score
                                    "foreground_properties": foreground_intensity_properties[k],
                                }
                                for k, v in dataset_json["channel_names"].items()
                            },
                            "num_classes": len(dataset_json["labels"].keys()),
                            "class_names": {
                                v: k for k, v in dataset_json["labels"].items()
                            },
                        },
                    }

                    json.dump(
                        config_dict,
                        f,
                        indent=4,
                    )

if __name__ == '__main__':
    export_onnx_model(
    dataset_name_or_id="001",
    output_dir="output_path",
    configurations={"2d"},
    batch_size=1)

2. 参数说明：

# 根据自己的设置传入参数
（1）dataset_name_or_id 填入 Dataset的id；
（2）output_dir：填入输出文件；
（3）configurations：训练的格式，例如2d，3d等；

3. 输出文件说明：

运行上面代码在输出文件夹下得到命名为：
“checkpoint_final.onnx 和 config.json” 文件；

checkpoint_final.onnx 为输出的onnx模型；