传输优化

减少模型输入输出数据在X86端和板端的传输量可提高工具的性能表现,工具可对如下三种使用场景提供传输优化支持。

  1. 模型输入固定或定期更新。若多帧推理模型的输入张量不变,那么输入张量不需要重复传输,只需在板端存储,后续推理直接复用板端张量。

  2. 模型输出过滤。未使用到的模型输出不需要回传至X86端。

  3. 多模型场景中前序模型输出可直接作为后续模型输入。此场景下,可将前序模型输出张量存储在板端且不回传,后续模型的对应输入亦无需传输,使用板端存储张量即可。

类、接口及参数

  1. HTensor

    HTensor用作传输优化场景中的输入输出张量,是X86端张量或板端张量的包装类,旨在提供统一的数据接口,并限制部分属性的修改以保证数据一致性。

    1). HTensor成员方法:__init__

    def __init__( self, data: Union[np.ndarray, torch.Tensor, None], device: Union[str, List[str], None], key: Optional[str] = None, ) -> None:

    初始化HTensor对象。

    • 参数
    参数说明
    data待封装的张量数据。
    device张量的存储设备,可选值包括 None"cpu""bpu"["cpu", "bpu"] 。其中 "cpu" 代表X86端, "bpu" 代表板端。
    key在板端用于唯一标识该张量的键值。

    2). HTensor属性:data

    获取或设置张量数据。设置数据时,若原数据不为 None,则新数据类型需要与原类型一致。

    3). HTensor属性:device

    获取张量存储设备信息。对象构造后禁止修改。

    4). HTensor属性:key

    获取张量在板端的唯一标识键值,在 device 中包含 "bpu" 时生效。对象构造后禁止修改。

    5). HTensor属性:shape

    获取张量形状。禁止修改,工具自动维护。

  2. HbmRpcSession成员方法 __call__ 中的 output_config 参数

    此参数用于配置当前推理帧结束后输出张量的传输行为,其类型为 Dict[str, Dict[str, Any]] ,一级键为模型输出名称,二级键需包含 "device""key" (可选)。其中 "device""key" 对应值的含义及约束与HTensor构造函数中 devicekey 参数保持一致。

    output_config 中正确配置了模型的某个输出名称时,在当前帧返回的推理结果中,对应输出张量将是HTensor类型。未配置的输出则会以一般类型( numpy.ndarraytorch.Tensor )返回。

传输优化应用示例

  1. 输入定期更新

    此示例中假设模型有一个名为 img 的输入,每 10 帧推理会更新一次,在 img 更新后的第一帧工具会将其传输至板端,其余帧不会有输入数据的传输。

    import torch import logging from hbm_infer.hbm_rpc_session import HbmRpcSession, HTensor, logger def periodic_input_update(epoch: int = 50): # 创建会话 session = HbmRpcSession(host=<available_ip>, local_hbm_path=<local_hbm_path>) # img 为定期更新,则需使用 HTensor 封装 fixed_img = HTensor( # 设置初始 data data=torch.ones((1, 3, 224, 224), dtype=torch.int8), # 固定输入或定期更新输入场景下,device 的设置只能为 ["cpu", "bpu"] # "cpu" 代表当前张量存储在 X86 端, "bpu" 代表将此张量将在板端创建存储副本 device=["cpu", "bpu"], # 由于 device 中包含 "bpu" ,故需要设置 key 用于在板端唯一标识此张量 key="model.input_0", ) # 创建模型 input dict input = {"img": fixed_img} for e in range(epoch): print(f"Epoch: {e}") # 满足一定条件时更新 img 输入 if e % 10 == 0: # 当 HTensor.data 更新后,在后续推理接口调用时,会把更新的张量数据传输至板端;未更新时,没有数据传输 fixed_img.data = torch.ones((1, 3, 224, 224), dtype=torch.int8) # 推理 output = session(data=input) for k, v in output.items(): print(k, v.shape) # 关闭会话 session.close_server() if __name__ == "__main__": logger.setLevel(logging.DEBUG) periodic_input_update(epoch=50)
  2. 输出过滤

    此示例中模型有三个输出,分别为 output_0output_1output_2 ,其中 output_2 无用被过滤掉,仅 output_0output_1 返回至X86端。

    import torch import logging from hbm_infer.hbm_rpc_session import HbmRpcSession, HTensor, logger def output_discard(): session = HbmRpcSession(host=<available_ip>, local_hbm_path=<local_hbm_path>) # 创建 input dict input = { "input_0": torch.ones((4, 1024, 1024), dtype=torch.float32), "input_1": torch.ones((4, 1024, 1024), dtype=torch.float32), } # 配置 output_config output_config = { # output_0 正常返回,故不配置 # output_1 正常返回,故不配置 # output_2 被过滤,即不回传至X86端,也不在板端做保存,设置 device 为 None 即可,由于 device 不包含 "bpu",故 key 不用设置 "output_2": {"device": None} } # 推理模型 output = session(data=input, output_config=output_config) # 在推理结果中,output_0 和 output_1 正常返回,类型为 torch.Tensor;output_2 类型则为 HTensor print(type(output["output_0"])) # <class 'torch.Tensor'> print(type(output["output_1"])) # <class 'torch.Tensor'> print(type(output["output_2"])) # <class 'hbm_infer.utils.HTensor'> # 其 data 属性为 None,即数据未回传 print(output["output_2"].data) # None # 无论输出是否被丢弃,其 shape 信息始终会回传 print(output["output_2"].shape) # (4, 1024, 1024) if __name__ == "__main__": logger.setLevel(logging.DEBUG) output_discard()
  3. 模型串联

    此示例假设传入的hbm文件中包含两个模型:model0与model1,其中model0名为 output_0 的输出将直接作为model1名为 input_0 的输入。此过程中,model0的输出无需回传至X86端,model1的输入也无需传输至板端。

    import torch import logging from hbm_infer.hbm_rpc_session import HbmRpcSession, HTensor, logger def model_chaining(): # 创建会话 session = HbmRpcSession(host=<available_ip>, local_hbm_path=<local_hbm_path>) # [model0, model1] # 注意:模型名称不等价于hbm文件名称 print(session.get_model_names()) # 创建 model0 的 input dict model0_input = {"input_0": torch.ones((4, 1024, 1024), dtype=torch.float32)} # 配置 model0 的 output_config model0_output_config = { "output_0": { # output_0 不用回传,直接在板端存储,故设置其 device 为 "bpu" 并设置 key "device": "bpu", "key": "model0.output_0", } } # 推理 model0 model0_output = session( data=model0_input, output_config=model0_output_config, model_name="model0" ) # model0.output_0 类型为 HTensor print(type(model0_output["output_0"])) # <class 'hbm_infer.utils.HTensor'> # 其 data 属性为 None,即数据未回传 print(model0_output["output_0"].data) # None # 无论输出张量是否被回传,其 shape 信息始终会回传 print(model0_output["output_0"].shape) # (4, 1024, 1024) # 创建 model1 的 input dict model1_input = { # 直接将 model0 的 output_0 推理结果作为 model1 的 input_0,后续推理时,model1 的 input_0 无需传输,而是根据 key 值直接复用板端张量 "input_0": model0_output["output_0"] } # 推理 model1,model1 结果正常返回,无需配置 output_config model1_output = session(data=model1_input, model_name="model1") print(type(model1_output["output_0"])) # <class 'torch.Tensor'> # 关闭会话 session.close_server() if __name__ == "__main__": logger.setLevel(logging.DEBUG) model_chaining()
  4. 综合应用

    下述推理pipeline覆盖输入定期更新、输出过滤及模型串联三个场景,流程图如下:

    hbm_infer_pipeline_sample

    参考代码如下:

    import torch import logging from hbm_infer.hbm_rpc_session import HbmRpcSession, HTensor, logger def test_pipeline(epoch=50): # 创建会话 session = HbmRpcSession(host=<available_ip>, local_hbm_path=<pipeline_hbm_path>) # [model0, model1, model2] # 注意:模型名称不等价于 hbm 文件名称 print(f"Model list: {session.get_model_names()}") # model1 的 input_1 为定期更新,需使用 HTensor 封装 model1_fixed_input1 = HTensor( # 设置初始 data data=torch.ones((4, 1024, 1024), dtype=torch.float32) * 2, # 固定输入或定期更新输入场景下,device 的设置只能为 ["cpu", "bpu"] # "cpu" 代表当前张量存储在 X86 端, "bpu" 代表将此张量将在板端创建存储副本 device=["cpu", "bpu"], # 由于 device 中包含 "bpu" ,故需要设置 key 用于在板端唯一标识此张量 key="model1.input_1", ) for e in range(epoch): print(f"Epoch: {e}") # model0 的 input_0 为正常输入 model0_input = { "input_0": torch.ones((4, 1024, 1024), dtype=torch.float32) * -1 } # 在板端保存 model0 的 output_0,且实际张量数据不回传至 X86 端,推理结果以 HTensor 形式返回,可直接作为后续模型的输入。 # 此 HTensor 的 data 属性将为 None,但您仍可通过 shape 属性获取其形状,作为后续模型的输入时,工具会自动根据其 key 属性在板端找到对应的张量。 model0_output_config = {"output_0": {"device": "bpu", "key": "model0.output_0"}} # 推理 model0 model0_output = session( data=model0_input, # 设置 output_config output_config=model0_output_config, # 指定要推理的模型名称 model_name="model0", ) # 模拟 model1 的 input_1 定期更新 if e % 10 == 0: model1_fixed_input1.data = ( torch.ones((4, 1024, 1024), dtype=torch.float32) * 2 ) model1_input = { # 直接使用 model0 的 output_0 作为 model1 的 input_0 "input_0": model0_output["output_0"], # Fixed input1 "input_1": model1_fixed_input1, } model1_output_config = { # 丢弃 output_0,设置 device 为 None 即可 "output_0": { "device": None, }, # 正常返回 output_1,所以不在 output_config 中配置 # 在板端存储 output_2 并且同时回传至 X86 端 "output_2": {"device": ["bpu", "cpu"], "key": "model1.output2"}, } # 推理 model1 model1_output = session( data=model1_input, output_config=model1_output_config, model_name="model1", ) # 使用 model1 的 output_2 作为 model2 的 input_0 model2_input = {"input_0": model1_output["output_2"]} # 推理 model2 model2_output = session(data=model2_input, model_name="model2") # 检查结果 print(f"{torch.all(model2_output['output_0'] == 1)}") # 关闭会话 session.close_server() if __name__ == "__main__": logger.setLevel(logging.DEBUG) test_pipeline(epoch=100)