更新时间:2024-09-10 GMT+08:00
分享

模型适配

MindSpore Lite是华为自研的推理引擎,能够最大化地利用昇腾芯片的性能。在使用MindSpore Lite进行离线推理时,需要先将模型转换为mindir模型,再利用MindSpore Lite作为推理引擎,将转换后的模型直接运行在昇腾设备上。模型转换需要使用converter_lite工具。

Huggingface提供的onnx模型文件的输入是动态shape,而mindir不支持动态shape,只能使用静态shape或者几个固定档位的分档shape代替。使用converter_lite转换模型时,也分为静态shape和分档shape两种方式,需要根据具体的业务需求使用对应的转换方式。本次迁移使用的是静态shape方式进行模型转换。

获取模型shape

由于在后续模型转换时需要知道待转换模型的shape信息,这里指导如何通过训练好的stable diffusion pytorch模型获取模型shape,主要有如下两种方式获取:

  • 方式一:通过stable diffusion的pytorch模型获取模型shape。
  • 方式二:通过查看ModelArts-Ascend代码仓库,根据每个模型的configs文件获取已知的shape大小。

下文主要介绍方式1如何通过stable diffusion的pytorch模型获取模型shape。

  1. 在pipeline应用准备章节,已经下载到sd的pytorch模型(/home_host/work/runwayml/pytorch_models)。进入工作目录:

    cd /home_host/work

  2. 新建python脚本文件“parse_models_shape.py”用于获取shape,其中model_path是指上面下载的pytorch_models的路径。

    # parse_models_shape.py
    import torch
    import numpy as np
    from diffusers import StableDiffusionPipeline
    
    model_path = '/home_host/work/runwayml/pytorch_models'
    pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32)
    
    # TEXT ENCODER
    num_tokens = pipeline.text_encoder.config.max_position_embeddings
    text_hidden_size = pipeline.text_encoder.config.hidden_size
    text_input = pipeline.tokenizer(
        "A sample prompt",
        padding="max_length",
        max_length=pipeline.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    print("# TEXT ENCODER")
    print(f"input_ids: {np.array(text_input.input_ids.shape).tolist()}")
    
    # UNET
    unet_in_channels = pipeline.unet.config.in_channels
    unet_sample_size = pipeline.unet.config.sample_size
    print("# UNET")
    print(f"sample: [{2}, {unet_in_channels} {unet_sample_size} {unet_sample_size}]")
    print(f"timestep: [{1}]")  # 此处应该是1,否则和后续的推理脚本不一致
    print(f"encoder_hidden_states: [{2}, {num_tokens} {text_hidden_size}]")
    
    # VAE ENCODER
    vae_encoder = pipeline.vae
    vae_in_channels = vae_encoder.config.in_channels
    vae_sample_size = vae_encoder.config.sample_size
    print("# VAE ENCODER")
    print(f"sample: [{1}, {vae_in_channels}, {vae_sample_size}, {vae_sample_size}]")
    
    # VAE DECODER
    vae_decoder = pipeline.vae
    vae_latent_channels = vae_decoder.config.latent_channels
    vae_out_channels = vae_decoder.config.out_channels
    print("# VAE DECODER")
    print(f"latent_sample: [{1}, {vae_latent_channels}, {unet_sample_size}, {unet_sample_size}]")
    
    # SAFETY CHECKER
    safety_checker = pipeline.safety_checker
    clip_num_channels = safety_checker.config.vision_config.num_channels
    clip_image_size = safety_checker.config.vision_config.image_size
    print("# SAFETY CHECKER")
    print(f"clip_input: [{1}, {clip_num_channels}, {clip_image_size}, {clip_image_size}]")
    print(f"images: [{1}, {vae_sample_size}, {vae_sample_size}, {vae_out_channels}]")

  3. 执行以下命令获取shape信息。

    python parse_models_shape.py

    可以看到获取的shape信息如下图所示。

    图1 shape信息

PyTorch模型转换为Onnx模型(可选)

获取onnx模型有两种方式,方式一是使用官方提供的模型转换脚本将pytorch模型转换为onnx模型,方式二是对于提供了onnx模型的仓库,可以直接下载onnx模型。下面介绍方式一如何操作,如果采用方式二,可以跳过此步骤。

  1. 通过git下载diffusers对应版本的源码。

    git clone https://github.com/huggingface/diffusers.git -b v0.11.1

  2. 在diffusers的script/convert_stable_diffusion_checkpoint_to_onnx.py脚本中,可以通过执行以下命令生成onnx模型,其中model_path指定pytorch的模型根目录,output_path指定生成的onnx模型目录。

    cd /home_host/work
    python diffusers/scripts/convert_stable_diffusion_checkpoint_to_onnx.py --model_path "./runwayml/pytorch_models"  --output_path "./pytorch_to_onnx_models"

静态shape模型转换

转换静态shape模型需要在模型转换阶段固定模型的输入shape,也就是说每个输入shape是唯一的。静态shape转换主要包括两种场景:

  • 第一种是待转换onnx模型的输入本身已经是静态shape,此时不需要在转换时指定输入shape也能够正常转换为和onnx模型输入shape一致的mindir模型。
  • 第二种是待转换onnx模型的输入是动态shape(导出onnx模型时指定了dynamic_axes参数),此时需要在转换时明确指定输入的shape。

转换时指定输入shape可以在命令行中指定,也可以通过配置文件的形式进行指定。

  • 在命令行中指定输入shape。

    命令行可以直接通过--inputShape参数指定输入的shape,格式为“input_name:input_shape”,如果有多个输入,需要使用“;”隔开,比如“input1_name:input1_shape;input2_name:input2_shape”。

    converter_lite --modelFile=./text_encoder/model.onnx --fmk=ONNX --saveType=MINDIR --optimize=ascend_oriented --outputFile=./text_encoder --inputShape="input_ids:1,77"
  • 在配置文件中指定输入shape。

    配置文件中通过“[ascend_context]”配置项指定input_shape,格式与命令行一致,多个输入,需要使用“;”隔开;然后在命令行中通过--configFile指定对应的配置文件路径即可。

    # text_encoder.ini 
    
    [ascend_context]
    input_shape=input_ids:[1,77]

    转换命令如下:

    converter_lite --modelFile=./text_encoder/model.onnx --fmk=ONNX --saveType=MINDIR --optimize=ascend_oriented --outputFile=./text_encoder --configFile=./text_encoder.ini

    在使用converter_lite工具转换时,默认是将所有算子的精度转换为fp16,如果想要将固定shape的模型精度修改为fp32进行转换,需要在配置文件中指定算子的精度模式为precision_mode,配置文件的写法如下(更多精度模式请参考precision_mode):

    # text_encoder.ini 
    [ascend_context]
    input_shape=input_ids:[1,77]
    precision_mode=enforce_fp32

对于本次AIGC迁移,为了方便对多个模型进行转换,可以通过批量模型转换脚本自动完成所有模型的转换。

  1. 执行以下命令创建并进入static_shape_convert目录。

    mkdir -p /home_host/work/static_shape_convert 
    cd /home_host/work/static_shape_convert

  2. 在static_shape_convert目录下新建converter_onnx2mindir.sh文件并复制下面内容。其中,onnx_dir表示onnx模型的目录,mindir_dir指定要生成的mindir模型的保存目录。

    # converter_onnx2mindir.sh
    # 设置onnx模型和mindir模型目录
    onnx_dir=/home_host/work/runwayml/onnx_models
    mindir_dir=./mindir_models
    
    # 指定配置文件路径
    config_dir=/home_host/work/modelarts-ascend/examples/AIGC/stable_diffusion/configs
    
    echo "================begin converter_lite====================="
    
    sub_cmd='--fmk=ONNX --optimize=ascend_oriented --saveType=MINDIR'
    mkdir -p $mindir_dir
    # rm缓存,慎改
    atc_data_dir=/root/atc_data/
    # 通用转换方法
    common_converter_model() {
        model_name=$1
        echo "start to convert $model_name"
        rm -rf $atc_data_dir
        converter_lite --modelFile="$onnx_dir/$model_name/model.onnx" \
                       --outputFile="$mindir_dir/$model_name" \
                       --configFile="$config_dir/$model_name.ini" \
                       $sub_cmd
        printf "end converter_lite\n"
    }
    common_converter_model "text_encoder"
    common_converter_model "unet"
    common_converter_model "vae_encoder"
    common_converter_model "vae_decoder"
    common_converter_model "safety_checker"
    
    echo "================converter_lite over====================="

    转换结果如下,其中safety_checker模型转换成功了,但中间有ERROR日志,该ERROR属于常量折叠失败,不影响结果。

    图2 转换结果

动态分档模型转换(可选)

如果迁移的模型有多个shape档位的需求,可以通过如下方式对模型进行分档转换。

动态分档是指将模型输入的某一维或者某几维设置为“动态”可变,但是需要提前设置可变维度的“档位”范围。即转换得到的模型能够在指定的动态轴上使用预设的几种shape(保证模型支持的shape),相比于静态shape更加灵活,且性能不会有劣化。

动态分档模型转换需要使用配置文件,指定输入格式为“ND”,并在config文件中配置ge.dynamicDims和input_shape使用,在input_shape中将输入shape的动态维度设为-1,并在ge.dynamicDims中指定动态维度的档位,更多配置项可以参考官方文档

  • 如果网络模型只有一个输入:每个档位的dim值与input_shape参数中的-1标识的参数依次对应,input_shape参数中有几个-1,则每档必须设置几个维度。
    以text_encoder模型为例,修改配置文件text_encoder.ini如下所示:
    # text_encoder.ini 
    
    [acl_build_options]
    input_format="ND"
    input_shape="input_ids:1,-1"
    ge.dynamicDims="77;33"

    使用上述配置文件转换得到的模型,支持的输入shape为(1,77)和(1,33)。

    然后使用converter lite执行模型转换,转换命令如下:

    converter_lite --modelFile=./onnx_models/text_encoder/model.onnx --fmk=ONNX --saveType=MINDIR --optimize=ascend_oriented --outputFile=./mindirs --configFile=./configs/text_encoder.ini
  • 如果网络模型有多个输入:档位的dim值与网络模型输入参数中的-1标识的参数依次对应,网络模型输入参数中有几个-1,则每档必须设置几个维度。
    以unet模型为例,该网络模型有三个输入,分别为“sample(1,4,64,64)”“timestep(1)”“encoder_hidden_states(1,77,768)”,修改unet.ini配置文件如下所示:
    # unet.ini 
    
    [acl_build_options]
    input_format="ND"
    input_shape="sample:-1,4,64,64;timestep:1;encoder_hidden_states:-1,77,768"
    ge.dynamicDims="1,1;2,2;3,3"

    转换得到的模型支持的输入dims组合档数分别为:

    图3 组合档数

    第0档:sample(1,4,64,64) + timestep(1) + encoder_hidden_states(1,77,768)

    第1档:sample(2,4,64,64) + timestep(1) + encoder_hidden_states(2,77,768)

    第2档:sample(3,4,64,64) + timestep(1) + encoder_hidden_states(3,77,768)

    然后使用converter lite执行模型转换,转换命令如下:

    converter_lite --modelFile=./onnx_models/unet/model.onnx --fmk=ONNX --saveType=MINDIR --optimize=ascend_oriented --outputFile=./mindirs --configFile=./configs/unet.ini
    • 最多支持100档配置,每一档通过英文逗号分隔。
    • 如果用户设置的dim数值过大或档位过多,可能会导致模型编译失败,此时建议用户减少档位或调低档位数值。
    • 如果用户设置了动态维度,实际推理时,使用的输入数据的shape需要与设置的档位相匹配。

相关文档