模型适配
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。
- 在pipeline应用准备章节,已经下载到sd的pytorch模型(/home_host/work/runwayml/pytorch_models)。进入工作目录:
cd /home_host/work
- 新建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}]")
- 执行以下命令获取shape信息。
python parse_models_shape.py
可以看到获取的shape信息如下图所示。
图1 shape信息
PyTorch模型转换为Onnx模型(可选)
获取onnx模型有两种方式,方式一是使用官方提供的模型转换脚本将pytorch模型转换为onnx模型,方式二是对于提供了onnx模型的仓库,可以直接下载onnx模型。下面介绍方式一如何操作,如果采用方式二,可以跳过此步骤。
- 通过git下载diffusers对应版本的源码。
git clone https://github.com/huggingface/diffusers.git -b v0.11.1
- 在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迁移,为了方便对多个模型进行转换,可以通过批量模型转换脚本自动完成所有模型的转换。
- 执行以下命令创建并进入static_shape_convert目录。
mkdir -p /home_host/work/static_shape_convert cd /home_host/work/static_shape_convert
- 在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需要与设置的档位相匹配。