내용 원문 및 사진 출처는 https://tvm.apache.org/docs/dev/convert_layout.html 입니다.

Background

data layout format이란 메모리에 어떻게 access 되는가를 나타냅니다. (Tensorflow는 NHWC layout을 지원합니다. 이는 locality에 많은 영향을 주기 때문에, performance에서 중요한 역할을 차지합니다.

TVM도 Framework parser, relay layout transformation, TOPI schedule과 같은 다양한 레벨에서 data layout format에 대해 신경써야 합니다. Relay에서 정의하는 ConvertLayout Pass에서는 layout handling에서 발생하는 다양한 복잡성들을 줄여줄 수 있습니다.

<aside> 💡 본문에는 motivation and overview와 design 파트도 적혀 있지만, 만약 tvm 코드를 이해하는 것이 목적이라면, 저처럼 바로 usage로 넘어가도 괜찮을 것 같습니다.

</aside>

Usage

ConvertLayout Pass는 relay.build pipeline의 기본 pass가 아닙니다. ConvertLayout Pass를 사용하려면, framework-to-relay parser와 relay.build module call 사이에서 호출해야 합니다.

우선, operator과 적용할 layouts의 매핑을 정의해야 합니다.

"default" 키워드를 사용하면, kernel layout을 TVM에서 지원하는 layout으로 자동으로 바꿔줍니다. 만약 user-specific-kernel-layout으로 하고 싶다면, 원하는 layout을 지정해 주면 됩니다.

# TFlite framework to Relay parser - **Default layout is NHWC**
mod, params = relay.frontend.from_tflite(tflite_model,
                                         shape_dict=shape_dict,
                                         dtype_dict=dtype_dict)

# We assume our model's heavily-layout sensitive operators only consist of nn.conv2d
**desired_layouts = {'nn.conv2d': ['NCHW', 'default']}**

# Convert the layout to NCHW
# RemoveUnunsedFunctions is used to clean up the graph.
seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
                                relay.transform.ConvertLayout(desired_layouts)])
with tvm.transform.PassContext(opt_level=3):
    mod = seq(mod)

# Call relay compilation
with relay.build_config(opt_level=3):
     graph, lib, params = relay.build(mod, target, params=params)
**desired_layouts = {'nn.conv2d': ['NCHW', 'OIHW']}**
pass = relay.transform.ConvertLayout(desired_layouts)

layout의 ordering은 register_convert_op_layout("OPNAME")의 implementation에 정의되어 있습니다. 위 예시에서는 [data_layout, kernel_layout]에 맞추어져 있습니다.

image classification에 필요한 웬만한 layout conversion들은 들어간 상태고, 그래프에 너무 많은 data layout transform이 들어가는 경우에는 별도의 작업이 필요합니다.