Convert Layout Pass
目录
1.背景
数据布局格式描述了如何在内存中放置数据。例如，卷积运算符的Tensorflow框架默认数据布局为NHWC，即数据为4维，并以行为主进行布局，其中N为第一维，C为最后一维。数据布局在模型性能中起主要作用，对空间和时间局部性有重大影响。例如，TVM中的Intel x86后端更喜欢使用NCHWc布局，其中C维被平铺为第2维，以有效利用数据局部性。同样，CUDA后端更喜欢将数据布局设置为NCHW格式。
本质上，TVM必须处理整个编译器工具链中的数据布局——框架解析器，Relay布局转换和TOPI调度。随着我们转向可能具有自己的数据布局限制的第三方代码生成集成，在TVM工具链中处理各个级别的布局将变得更具挑战性。因此，我们开发了新的Relay转换-ConvertLayout –以减少由于布局处理而引起的一些复杂性。
如果您直接想了解ConvertLayout Pass的用法，请直接跳至第4部分-用法。
2.动机和概述
让我们看一个简单的场景，以了解由于布局不同而引起的复杂性——假设我们要为ARM边缘设备编译Tensorflow NHWC图。但是，假设我们目前在TOPI 对ARM仅支持NCHW调度。因此，框架布局和TOPI支持的布局之间不匹配。解决这种不匹配问题的一种方法是在每次卷积之前和之后插入布局转换，以使最终的卷积具有NCHW输入数据布局并可以使用TOPI调度。但是，由于存在太多的布局转换，这可能导致性能下降。
我们在其他用例中也遇到了类似的问题
-
无法在Nvidia GPU上运行TFLite图形。TOPI针对GPU仅拥有NCHW调度。
-
【AlterOpLayout】中不断复杂的逻辑，用于卷积以支持不同的布局转换对。
-
由于额外的布局转换，TF图的性能欠佳。
-
TensorRT等第三方代码生成集成的复杂性，它更喜欢数据布局采用一种格式。
为了解决这些问题，我们引入了ConvertLayout转换，该转换设置了基础结构，以最少的数据布局转换次数即可更改整个图形的数据布局。在理想情况下，我们将只对数据进行2个布局转换，一个在开始，一个在结尾。下面是显示转换的示例
# Original graph - 2 convolutions in NHWC format.
fn (%x: Tensor[(1, 56, 56, 64), float32], %weight1: Tensor[(3, 3, 64, 32), float32], %weight2: Tensor[(3, 3, 32, 32), float32]) {
%0 = nn.conv2d(%x, %weight1, padding=[1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO");
%1 = nn.relu(%0);
%2 = nn.conv2d(%1, %weight2, padding=[1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO");
nn.relu(%2)
}
# After ConvertLayout - For data, there is a transform at the start and at the end.
# For weights, there are transforms to adapt to NCHW layout. These will be removed by FoldConstant pass.
fn (%x: Tensor[(1, 56, 56, 64), float32], %weight1: Tensor[(3, 3, 64, 32), float32], %weight2: Tensor[(3, 3, 32, 32), float32]) {
%0 = layout_transform(%x, src_layout="NHWC", dst_layout="NCHW") /* ty=Tensor[(1, 64, 56, 56), float32] */;
%1 = layout_transform(%weight1, src_layout="HWIO", dst_layout="OIHW") /* ty=Tensor[(32, 64, 3, 3), float32] */;
%2 = nn.conv2d(%0, %1, padding=[1, 1], channels=32, kernel_size=[3, 3]) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%3 = nn.relu(%2) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%4 = layout_transform(%weight2, src_layout="HWIO", dst_layout="OIHW") /* ty=Tensor[(32, 32, 3, 3), float32] */;
%5 = nn.conv2d(%3, %4, padding=[1, 1], channels=32, kernel_size=[3, 3]) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%6 = nn.relu(%5) /* ty=Tensor[(1, 32, 56, 56), float32] */;
layout_transform(%6, src_layout="NCHW", dst_layout="NHWC") /* ty=Tensor[(1, 56, 56, 32), float32] */
}
3.设计
在研究ConvertLayout转换之前，让我们根据运算符对数据布局的敏感性将其分为3类。此分类将在以后理解Convertlayout转换详细信息时很有用。
-
Layout agnostic -Relu，pow等。这些运算符不受数据布局的影响，无论功能还是性能。
-
Lightly-layout sensitive -pad, concatenate, 像sum这样的reduce操作等。这些运算符具有某些属性，如果我们在其之前进行布局转换，这些属性会在功能上受到影响。但是，就性能而言，差异并不明显。对于这些操作符，仅适应先前的操作符输出数据布局是有益的。
-
Heavily-layout sensitive -卷积，conv2d_transpose等。这些运算符在功能和性能方面都受到数据布局的严重影响。它们还把数据布局作为操作符属性。通常，为这些运算符修改输入数据布局是有益的（如果它不是高效的数据布局），而其余与布局无关的和轻度布局敏感的运算符将适应由这些重布局敏感的输出所控制的布局操作符。
现在让我们看一下两个相关的Relay操作符属性。每个Relay操作符都有可以由TVM开发人员定义的属性，例如InferType。通常，Relay转换逐个操作的遍历图，并读取这些操作符的属性。例如，InferType转换查看运算符的【InferType】属性，确定其输出形状和类型，然后将其传递给下一个运算符InferType属性。同样，在我们的上下文中，我们有2个此类属性-【FTVMConvertLayout】和【FInferCorrectLayout】。【ConvertLayout】遍历图形并查看这2个属性以及一个自动布局转换插入模块以处理数据布局。因此，整个过程可以分为三个步骤：
-
运行【FTVMConvertLayout】属性-这使开发人员可以将原始Relay 【expr】转换为具有新布局的新Relay 【expr】，从而允许用户定义布局更改。为了方便开发人员，有一个python回调。仅用于布局严重敏感的运算符。
-
运行【FTVMInferCorretLayout】属性-我们可以将其视为布局推断。它查看原始输入布局和新输入布局，这些布局来自先前的运算符或来自【FTVMConvertLayout】修改的【expr】（如果已使用）。 轻度布局的敏感运算符可以使用此属性，以使其属性适应新的数据布局。布局推断发生在每个运算符上。
-
自动插入布局转换-上一步布局推断，为输入表达式设置新布局。如果这些布局与原始布局不同，则此组件将自动插入一个布局转换。因此，开发人员无需为此组件做任何事情。
这些步骤按顺序对每个运算符进行，其中ConvertLayout转换继续将新布局传递给下一个运算符属性，最终导致逐个运算符修改整个图。现在，让我们看一下如何定义两个属性的一对示例。
FTVMConvertLayout-用于更改布局的Python回调 -用于严重布局敏感的运算符。例如，可以返回具有新数据和内核布局的新卷积运算符。其他2个组件将推断布局并在需要时插入布局转换。卷积运算符的一个示例如下所示，其中我们正在转换为NCHW布局。
@reg.register_convert_op_layout("nn.conv2d")
def convert_conv2d(attrs, inputs, tinfos, desired_layout):
"""Convert Layout pass registration for conv2d op.
Parameters
----------
attrs : tvm.attrs.Attrs
Attributes of current convolution
inputs : list of tvm.relay.Expr
The args of the Relay expr to be legalized
tinfos : list of types
List of input and output types
desired_layout : str
The desired layout
Returns
-------
result : tvm.relay.Expr
The transformed expr
"""
from tvm import relay
data_layout = attrs['data_layout']
kernel_layout = attrs['kernel_layout']
data, weight = inputs
assert desired_layout == 'NCHW', \
"Currently only transformation to NCHW layout is supported."
if desired_layout == 'NCHW':
new_attrs = dict(attrs)
new_attrs['data_layout'] = desired_layout
new_attrs['kernel_layout'] = 'OIHW'
# Actual insertion of layout transforms is taken care internally
# by ConvertLayout pass.
return relay.nn.conv2d(data, weight, **new_attrs)
return None
FInferCorrectLayout-布局推断 -当前，此属性仅在C ++中公开。此函数采用原始输入布局和新的输入布局（从上一个运算符或python回调转换以进行布局更改），并推断最终的数据布局。为每个运算符调用布局推断。对于不同的操作符类别，用法可能有所不同。对于与布局无关的运算符，我们只想在此函数中返回新的数据布局。对于轻度布局和重度布局敏感的运算符，我们可以更改运算符属性（例如，concatenate操作的axis，pad操作的pad_width），以便我们可以适应新的数据布局，从而避免插入布局变换。让我们看几个例子，以更好地理解这一点。
第一个示例是与布局无关的运算符。这些运算符没有任何受数据布局影响的运算符属性，因此我们仅适应新的布局。
// For operator set its attributes like following
// .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout);
// Take arbitrary input layouts and copy to outputs.
inline Array<Array<Layout> > ElemwiseArbitraryLayout(const Attrs& attrs,
const Array<Layout>& new_in_layouts,
const Array<Layout>& old_in_layouts,
const Array<Array<IndexExpr>> &old_in_shapes) {
Layout ret;
if (new_in_layouts.defined()) {
CHECK_GE(new_in_layouts.size(), 1);
ret = new_in_layouts[0];
} else {
for (size_t i = 0; i < old_in_layouts.size(); ++i) {
if (old_in_layouts[i].defined()) {
ret = old_in_layouts[i];
break;
}
}
}
return Array<Array<Layout> >{Array<Layout>(old_in_layouts.size(), ret), {ret}};
}
第二个示例是一个轻微布局敏感的运算符-批量归一化。BatchNorm有一个axis运算符，当我们从NHWC转到NCHW数据布局时，必须更改它。（对于重度布局敏感操作符也需要类似的处理）
Array<Array<Layout>> BatchNormInferCorrectLayout(const Attrs& attrs,
const Array<Layout>& new_in_layouts,
const Array<Layout>& old_in_layouts,
const Array<Array<IndexExpr>>& old_in_shapes) {
BatchNormAttrs* param = const_cast<BatchNormAttrs*>(attrs.as<BatchNormAttrs>());
size_t axis =
param->axis < 0 ? param->axis + old_in_shapes[0].size() : static_cast<size_t>(param->axis);
Layout ret = Layout::Undef();
// For example, consider old_layout = NHWC, and new_layout = NCHW, and param->axis = 3
if (new_in_layouts.defined() && old_in_layouts.defined()) {
// Get the new C axis. Extract the dim in old layout. Find the index of that dim in next layout.
// Following line gives bn_dim = C as old_layout = NHWC, axis = 3
const auto& bn_dim = old_in_layouts[0][axis];
// The new_index is 1 because new_layout = NCHW and bn_dim is C
auto new_index = new_in_layouts[0].IndexOf(bn_dim);
// We modify the layout-dependent attribute here - axis to 1.
param->axis = new_index;
// Finally, we adapt to the new layout.
ret = new_in_layouts[0];
} else if (old_in_layouts.defined()) {
ret = old_in_layouts[0];
}
// In case both new and old layouts are undefined, then there is no need of a change.
// ConvertLayout pass skips the automatic insertion of layout transforms in this case.
// Following line is not important to tutorial. But, layout inference needs to define
// the layout for all input and output data layouts. For batch norm, the other inputs
// and outputs are vector having length of C dim in the input. So, we set the other
// layouts as C. BN has 5 inputs, 3 outputs. The last 4 inputs and last 2 outputs
// have "C" layout.
Layout c_layout = Layout("C");
return Array<Array<Layout>>{{ret, c_layout, c_layout, c_layout, c_layout},
{ret, c_layout, c_layout}};
}
4.用法
ConvertLayout转换非常容易使用。该转换不是默认的relay.build管道的一部分。预期的用途是在framework-to-relay解析器和relay.build模块调用之间调用它。
# 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)
# Convert the layout to NCHW
# RemoveUnunsedFunctions is used to clean up the graph.
seq = relay.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout('NCHW')])
with relay.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)
当前实现支持几乎所有图像分类模型中常用的运算符。但是，如果在图形中遇到太多的数据布局转换，则很可能有一个运算符需要对其布局进行特殊处理，如第3节中所述。一些在这种情况下可以提供帮助的拉取请求是
扫一扫
1034
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
