三种版本的 DFL#
import set_env
from d2py.utils.file import mkdir
root_dir = ".temp"
mkdir(f"{root_dir}/logs")
import torch
from torch.nn import functional as F
from torch import nn
from torch.onnx import OperatorExportTypes, utils
class DFLV1(nn.Module):
"""
Integral module of Distribution Focal Loss (DFL).
Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
"""
def __init__(self, c1=16):
"""Initialize a convolutional layer with a given number of input channels."""
super().__init__()
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
self.c1 = c1
def forward(self, x):
"""Applies a transformer layer on input tensor 'x' and returns a tensor."""
b, c, a = x.shape # batch, channels, anchors
x = x.view(b, 4, self.c1, a)
x = x.transpose(3, 1).transpose(2, 3)
x = x.softmax(3)
x = x.transpose(3, 1)
x = self.conv(x)
return x.view(b, 4, a)
class DFLV2(nn.Module):
"""
Integral module of Distribution Focal Loss (DFL).
Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
"""
def __init__(self, c1=16):
"""Initialize a convolutional layer with a given number of input channels."""
super().__init__()
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
self.c1 = c1
def forward(self, x):
"""Applies a transformer layer on input tensor 'x' and returns a tensor."""
b, c, a = x.shape # batch, channels, anchors
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
class DFLV3(nn.Module):
"""
Integral module of Distribution Focal Loss (DFL).
Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
"""
def __init__(self, c1=16):
"""Initialize a convolutional layer with a given number of input channels."""
super().__init__()
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
self.c1 = c1
def forward(self, x):
"""Applies a transformer layer on input tensor 'x' and returns a tensor."""
b, c, a = x.shape # batch, channels, anchors
return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
class DFL(nn.Module):
def __init__(self, c1=16):
super().__init__()
self.conv0 = nn.Conv2d(3, 16, 1, bias=False)
self.conv1 = nn.Conv2d(16, 64, 1, bias=False)
self.conv2 = nn.Conv2d(16, 64, 1, bias=False)
self.conv3 = nn.Conv2d(16, 100, 1, bias=False)
self.v1 = DFLV1(c1)
self.v2 = DFLV2(c1)
self.v3 = DFLV3(c1)
def forward(self, x):
x = self.conv0(x)
x1 = self.conv1(x).view(1, 64, -1)
x2 = self.conv2(x).view(1, 64, -1)
x3 = self.conv3(x).view(1, 64, -1)
x1 = self.v1(x1)
x2 = self.v2(x2)
x3 = self.v3(x3)
return x1, x2, x3
model = DFL().eval()
shape = 1, 3, 48, 80
xx = torch.rand(*shape, dtype=torch.float32, requires_grad=False)
# model = torch.jit.trace(model, xx)
# 导出模型
input_name = "data"
output_name = "dfl"
utils.export(
model, # torch 模型
xx, # 模型输入或者对于多个输入，使用元组
f"{root_dir}/{output_name}.onnx", # 模型保存的位置（可以是文件或类似文件的对象）
export_params=True, # 将训练后的参数权重存储在模型文件内
opset_version=17, # 导出模型的 ONNX 版本
do_constant_folding=True, # 是否执行常量折叠以进行优化
input_names = [input_name], # 模型的输入名称
output_names = ['output'], # 模型的输出名称
keep_initializers_as_inputs=True,
# export_modules_as_functions=True,
verbose=True,
operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH,
# dynamic_axes={'data' : {0 : 'batch_size'}, # 可变长度的轴
# 'output' : {0 : 'batch_size'}}
)
Exported graph: graph(%data : Float(1, 3, 48, 80, strides=[11520, 3840, 80, 1], requires_grad=0, device=cpu),
%conv0.weight : Float(16, 3, 1, 1, strides=[3, 1, 1, 1], requires_grad=1, device=cpu),
%conv1.weight : Float(64, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=1, device=cpu),
%conv2.weight : Float(64, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=1, device=cpu),
%conv3.weight : Float(100, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=1, device=cpu),
%v1.conv.weight : Float(1, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=0, device=cpu),
%v2.conv.weight : Float(1, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=0, device=cpu),
%v3.conv.weight : Float(1, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=0, device=cpu)):
%/conv0/Conv_output_0 : Float(1, 16, 48, 80, strides=[61440, 3840, 80, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/conv0/Conv"](%data, %conv0.weight), scope: __main__.DFL::/torch.nn.modules.conv.Conv2d::conv0 # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/conv1/Conv_output_0 : Float(1, 64, 48, 80, strides=[245760, 3840, 80, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/conv1/Conv"](%/conv0/Conv_output_0, %conv1.weight), scope: __main__.DFL::/torch.nn.modules.conv.Conv2d::conv1 # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/Constant_output_0 : Long(3, strides=[1], device=cpu) = onnx::Constant[value= 1 64 -1 [ CPULongType{3} ], onnx_name="/Constant"](), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:82:0
%/Reshape_output_0 : Float(1, 64, 3840, strides=[245760, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/Reshape"](%/conv1/Conv_output_0, %/Constant_output_0), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:82:0
%/conv2/Conv_output_0 : Float(1, 64, 48, 80, strides=[245760, 3840, 80, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/conv2/Conv"](%/conv0/Conv_output_0, %conv2.weight), scope: __main__.DFL::/torch.nn.modules.conv.Conv2d::conv2 # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/Constant_1_output_0 : Long(3, strides=[1], device=cpu) = onnx::Constant[value= 1 64 -1 [ CPULongType{3} ], onnx_name="/Constant_1"](), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:83:0
%/Reshape_1_output_0 : Float(1, 64, 3840, strides=[245760, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/Reshape_1"](%/conv2/Conv_output_0, %/Constant_1_output_0), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:83:0
%/conv3/Conv_output_0 : Float(1, 100, 48, 80, strides=[384000, 3840, 80, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/conv3/Conv"](%/conv0/Conv_output_0, %conv3.weight), scope: __main__.DFL::/torch.nn.modules.conv.Conv2d::conv3 # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/Constant_2_output_0 : Long(3, strides=[1], device=cpu) = onnx::Constant[value= 1 64 -1 [ CPULongType{3} ], onnx_name="/Constant_2"](), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:84:0
%/Reshape_2_output_0 : Float(1, 64, 6000, strides=[384000, 6000, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/Reshape_2"](%/conv3/Conv_output_0, %/Constant_2_output_0), scope: __main__.DFL:: # /tmp/ipykernel_3296793/673152664.py:84:0
%/v1/Constant_output_0 : Long(4, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 4 16 3840 [ CPULongType{4} ], onnx_name="/v1/Constant"](), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:24:0
%/v1/Reshape_output_0 : Float(1, 4, 16, 3840, strides=[245760, 61440, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v1/Reshape"](%/Reshape_output_0, %/v1/Constant_output_0), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:24:0
%/v1/Transpose_output_0 : Float(1, 3840, 4, 16, strides=[245760, 1, 61440, 3840], requires_grad=1, device=cpu) = onnx::Transpose[perm=[0, 3, 1, 2], onnx_name="/v1/Transpose"](%/v1/Reshape_output_0), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:25:0
%/v1/Softmax_output_0 : Float(1, 3840, 4, 16, strides=[245760, 64, 16, 1], requires_grad=1, device=cpu) = onnx::Softmax[axis=3, onnx_name="/v1/Softmax"](%/v1/Transpose_output_0), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:26:0
%/v1/Transpose_1_output_0 : Float(1, 16, 4, 3840, strides=[245760, 1, 16, 64], requires_grad=1, device=cpu) = onnx::Transpose[perm=[0, 3, 2, 1], onnx_name="/v1/Transpose_1"](%/v1/Softmax_output_0), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:27:0
%/v1/conv/Conv_output_0 : Float(1, 1, 4, 3840, strides=[15360, 15360, 3840, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/v1/conv/Conv"](%/v1/Transpose_1_output_0, %v1.conv.weight), scope: __main__.DFL::/__main__.DFLV1::v1/torch.nn.modules.conv.Conv2d::conv # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/v1/Constant_1_output_0 : Long(3, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 4 3840 [ CPULongType{3} ], onnx_name="/v1/Constant_1"](), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:29:0
%output : Float(1, 4, 3840, strides=[15360, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v1/Reshape_1"](%/v1/conv/Conv_output_0, %/v1/Constant_1_output_0), scope: __main__.DFL::/__main__.DFLV1::v1 # /tmp/ipykernel_3296793/673152664.py:29:0
%/v2/Constant_output_0 : Long(4, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 4 16 3840 [ CPULongType{4} ], onnx_name="/v2/Constant"](), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%/v2/Reshape_output_0 : Float(1, 4, 16, 3840, strides=[245760, 61440, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v2/Reshape"](%/Reshape_1_output_0, %/v2/Constant_output_0), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%/v2/Transpose_output_0 : Float(1, 16, 4, 3840, strides=[245760, 3840, 61440, 1], requires_grad=1, device=cpu) = onnx::Transpose[perm=[0, 2, 1, 3], onnx_name="/v2/Transpose"](%/v2/Reshape_output_0), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%/v2/Softmax_output_0 : Float(1, 16, 4, 3840, strides=[245760, 15360, 3840, 1], requires_grad=1, device=cpu) = onnx::Softmax[axis=1, onnx_name="/v2/Softmax"](%/v2/Transpose_output_0), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%/v2/conv/Conv_output_0 : Float(1, 1, 4, 3840, strides=[15360, 15360, 3840, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/v2/conv/Conv"](%/v2/Softmax_output_0, %v2.conv.weight), scope: __main__.DFL::/__main__.DFLV2::v2/torch.nn.modules.conv.Conv2d::conv # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/v2/Constant_1_output_0 : Long(3, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 4 3840 [ CPULongType{3} ], onnx_name="/v2/Constant_1"](), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%66 : Float(1, 4, 3840, strides=[15360, 3840, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v2/Reshape_1"](%/v2/conv/Conv_output_0, %/v2/Constant_1_output_0), scope: __main__.DFL::/__main__.DFLV2::v2 # /tmp/ipykernel_3296793/673152664.py:48:0
%/v3/Constant_output_0 : Long(4, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 16 4 6000 [ CPULongType{4} ], onnx_name="/v3/Constant"](), scope: __main__.DFL::/__main__.DFLV3::v3 # /tmp/ipykernel_3296793/673152664.py:67:0
%/v3/Reshape_output_0 : Float(1, 16, 4, 6000, strides=[384000, 24000, 6000, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v3/Reshape"](%/Reshape_2_output_0, %/v3/Constant_output_0), scope: __main__.DFL::/__main__.DFLV3::v3 # /tmp/ipykernel_3296793/673152664.py:67:0
%/v3/Softmax_output_0 : Float(1, 16, 4, 6000, strides=[384000, 24000, 6000, 1], requires_grad=1, device=cpu) = onnx::Softmax[axis=1, onnx_name="/v3/Softmax"](%/v3/Reshape_output_0), scope: __main__.DFL::/__main__.DFLV3::v3 # /tmp/ipykernel_3296793/673152664.py:67:0
%/v3/conv/Conv_output_0 : Float(1, 1, 4, 6000, strides=[24000, 24000, 6000, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name="/v3/conv/Conv"](%/v3/Softmax_output_0, %v3.conv.weight), scope: __main__.DFL::/__main__.DFLV3::v3/torch.nn.modules.conv.Conv2d::conv # /media/pc/data/tmp/cache/conda/envs/py312x/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0
%/v3/Constant_1_output_0 : Long(3, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1 4 6000 [ CPULongType{3} ], onnx_name="/v3/Constant_1"](), scope: __main__.DFL::/__main__.DFLV3::v3 # /tmp/ipykernel_3296793/673152664.py:67:0
%88 : Float(1, 4, 6000, strides=[24000, 6000, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/v3/Reshape_1"](%/v3/conv/Conv_output_0, %/v3/Constant_1_output_0), scope: __main__.DFL::/__main__.DFLV3::v3 # /tmp/ipykernel_3296793/673152664.py:67:0
return (%output, %66, %88)
from copy import deepcopy
import onnx
import tvm
from tvm import relay
onnx_model = onnx.load(f"{root_dir}/{output_name}.onnx")
mod, params = relay.frontend.from_onnx(onnx_model, {"data": shape}, freeze_params=True)
mod = relay.transform.InferType()(mod)
origin_mod = deepcopy(mod)
# with tvm.transform.PassContext(opt_level=3):
# mod = relay.quantize.prerequisite_optimize(mod, params)
mod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), float32] span=/conv0/Conv.conv0.weight:0:0 */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1]) /* ty=Tensor[(1, 16, 48, 80), float32] span=/conv0/Conv:0:0 */;
%1 = nn.conv2d(%0, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv1/Conv.conv1.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv1/Conv:0:0 */;
%2 = reshape(%1, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape:0:0 */;
%3 = reshape(%2, newshape=[1, 4, 16, 3840]) /* ty=Tensor[(1, 4, 16, 3840), float32] span=/v1/Reshape:0:0 */;
%4 = transpose(%3, axes=[0, 3, 1, 2]) /* ty=Tensor[(1, 3840, 4, 16), float32] span=/v1/Transpose:0:0 */;
%5 = nn.softmax(%4, axis=3) /* ty=Tensor[(1, 3840, 4, 16), float32] span=/v1/Softmax:0:0 */;
%6 = transpose(%5, axes=[0, 3, 2, 1]) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v1/Transpose_1:0:0 */;
%7 = nn.conv2d(%6, meta[relay.Constant][2] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v1/conv/Conv.v1.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 3840), float32] span=/v1/conv/Conv:0:0 */;
%8 = nn.conv2d(%0, meta[relay.Constant][3] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv2/Conv.conv2.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv2/Conv:0:0 */;
%9 = reshape(%8, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape_1:0:0 */;
%10 = reshape(%9, newshape=[1, 4, 16, 3840]) /* ty=Tensor[(1, 4, 16, 3840), float32] span=/v2/Reshape:0:0 */;
%11 = transpose(%10, axes=[0, 2, 1, 3]) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v2/Transpose:0:0 */;
%12 = nn.softmax(%11, axis=1) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v2/Softmax:0:0 */;
%13 = nn.conv2d(%12, meta[relay.Constant][4] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v2/conv/Conv.v2.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 3840), float32] span=/v2/conv/Conv:0:0 */;
%14 = nn.conv2d(%0, meta[relay.Constant][5] /* ty=Tensor[(100, 16, 1, 1), float32] span=/conv3/Conv.conv3.weight:0:0 */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1]) /* ty=Tensor[(1, 100, 48, 80), float32] span=/conv3/Conv:0:0 */;
%15 = reshape(%14, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 6000), float32] span=/Reshape_2:0:0 */;
%16 = reshape(%15, newshape=[1, 16, 4, 6000]) /* ty=Tensor[(1, 16, 4, 6000), float32] span=/v3/Reshape:0:0 */;
%17 = nn.softmax(%16, axis=1) /* ty=Tensor[(1, 16, 4, 6000), float32] span=/v3/Softmax:0:0 */;
%18 = nn.conv2d(%17, meta[relay.Constant][6] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v3/conv/Conv.v3.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 6000), float32] span=/v3/conv/Conv:0:0 */;
%19 = reshape(%7, newshape=[1, 4, 3840]) /* ty=Tensor[(1, 4, 3840), float32] span=/v1/Reshape_1:0:0 */;
%20 = reshape(%13, newshape=[1, 4, 3840]) /* ty=Tensor[(1, 4, 3840), float32] span=/v2/Reshape_1:0:0 */;
%21 = reshape(%18, newshape=[1, 4, 6000]) /* ty=Tensor[(1, 4, 6000), float32] span=/v3/Reshape_1:0:0 */;
(%19, %20, %21) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
注册算子 vta_special.yolo_dfl#
import numpy as np
from tvm.relay.testing import run_infer_type
from tvm.relay.dataflow_pattern import (
wildcard, is_op,
is_constant,
DFPatternCallback,
rewrite
)
import tvm
from tvm.ir.attrs import DictAttrs
from tvm.relay import transform as _transform
from tvm import relay, te, topi
from tvm.relay.op import op as _op
from tvm.target import generic_func
@generic_func
def schedule_special_op(attrs, outs, target):
with target:
outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
output = outs[0]
sch = te.create_schedule(output.op)
return sch
def custom_yolo_dfl_rel(arg_types, attrs):
assert len(arg_types) == 1, "type relation arg number mismatch!"
if attrs:
assert isinstance(attrs, DictAttrs)
in_shape = attrs.in_shape
bbox_size = 4
assert in_shape[1]%bbox_size == 0
out_shape = (in_shape[0], bbox_size, in_shape[2])
return relay.TensorType(out_shape, "float32")
op_name = "vta_special.yolo_dfl"
_op.register(op_name, r"code(cal yolo_dfl.)code")
_op.get(op_name).set_num_inputs(1)
_op.get(op_name).add_argument("data", "Tensor", "The input data tensor.")
_op.get(op_name).set_attrs_type_key("DictAttrs")
_op.get(op_name).add_type_rel(op_name, custom_yolo_dfl_rel)
_op.get(op_name).set_support_level(1)
_op.register_pattern(op_name, _op.OpPattern.COMM_REDUCE)
_op.register_stateful(op_name, False) # 无状态算子
def yolo_dfl(x, channel, in_shape, version="v3", x_scale=-1, x_split=-1):
attrs = tvm.ir.make_node(
"DictAttrs",
channel=channel, in_shape=in_shape, version=version,
x_scale=x_scale, x_split=x_split,
)
return relay.Call(_op.get(op_name), [x], attrs=attrs, type_args=None, span=None)
@_op.register_compute(op_name)
def output_yolo_dfl_compute(attrs, inputs, out_type):
"""yolo_dfl Relay 计算"""
assert len(inputs) == 1, "输入参数数量不为 1"
x = inputs[0]
b, c, a = attrs.in_shape # batch, channels, anchors
assert c % 4 == 0
if x.dtype == "int8":
x = topi.cast(x, "float32")
x = topi.multiply(x, attrs.x_scale)
w = topi.arange(0, attrs.channel, dtype="float32")
w = topi.reshape(w, (1, attrs.channel, 1, 1))
if attrs.version == "v3":
x = topi.reshape(x, (b, c//4, 4, a))
x = topi.nn.softmax(x, axis=1)
x = topi.nn.conv2d(x, w, padding=[0, 0, 0, 0], strides=(1, 1), dilation=(1, 1))
x = topi.reshape(x, (b, 4, a))
elif attrs.version == "v2":
x = topi.reshape(x, (b, 4, c//4, a))
x = topi.transpose(x, [0, 2, 1, 3])
x = topi.nn.softmax(x, axis=1)
x = topi.nn.conv2d(x, w, padding=[0, 0, 0, 0], strides=(1, 1), dilation=(1, 1))
x = topi.reshape(x, (b, 4, a))
elif attrs.version == "v1":
x = topi.reshape(x, (b, 4, c//4, a))
x = topi.transpose(x, [0, 3, 1, 2])
x = topi.nn.softmax(x, axis=3)
x = topi.transpose(x, [0, 3, 2, 1])
x = topi.nn.conv2d(x, w, padding=[0, 0, 0, 0], strides=(1, 1), dilation=(1, 1))
x = topi.reshape(x, (b, 4, a))
else:
raise TypeError(f"暂未支持 {attrs.version}")
return [x]
_op.register_schedule(op_name, schedule_special_op) # 定义调度
GenericFunc(0x9c088c0)
class DFLV1Rewrite(DFPatternCallback):
def __init__(self):
super().__init__()
self.x = wildcard()
self.reshape = is_op("reshape")(self.x)
self.transpose = is_op("transpose")(self.reshape).has_attr({"axes": [0, 3, 1, 2]})
self.softmax = is_op("nn.softmax")(self.transpose).has_attr({"axis": 3})
self.transpose2 = is_op("transpose")(self.softmax).has_attr({"axes": [0, 3, 2, 1]})
self.conv_weight = is_constant()
self.conv = is_op("nn.conv2d")(self.transpose2, self.conv_weight)
self.reshape2 = is_op("reshape")(self.conv)
self.pattern = self.reshape2
def callback(self, pre, post, node_map):
x = node_map[self.x][0]
conv_weight = node_map[self.conv_weight][0]
# conv = node_map[self.conv][0]
b, c, a = _transform.InferTypeLocal(x).shape # batch, channels, anchors
conv_weight_shape = _transform.InferTypeLocal(conv_weight).shape
# print(f"conv_weight_shape[2]: {type(conv_weight_shape[2])}")
# print(dict(conv.attrs))
assert conv_weight_shape[0] == conv_weight_shape[2] == conv_weight_shape[3] == 1
# x = yolo_dfl(x, int(conv_weight_shape[1]), (b, c, a))
x = yolo_dfl(x, conv_weight_shape[1], (b, c, a), version="v1")
return x
class DFLV2Rewrite(DFPatternCallback):
def __init__(self):
super().__init__()
self.x = wildcard()
self.reshape = is_op("reshape")(self.x)
self.transpose = is_op("transpose")(self.reshape).has_attr({"axes": [0, 2, 1, 3]})
self.softmax = is_op("nn.softmax")(self.transpose).has_attr({"axis": 1})
self.conv_weight = is_constant()
self.conv = is_op("nn.conv2d")(self.softmax, self.conv_weight)
self.reshape2 = is_op("reshape")(self.conv)
self.pattern = self.reshape2
def callback(self, pre, post, node_map):
x = node_map[self.x][0]
conv_weight = node_map[self.conv_weight][0]
b, c, a = _transform.InferTypeLocal(x).shape # batch, channels, anchors
conv_weight_shape = _transform.InferTypeLocal(conv_weight).shape
assert conv_weight_shape[0] == conv_weight_shape[2] == conv_weight_shape[3] == 1
x = yolo_dfl(x, conv_weight_shape[1], (b, c, a), version="v2")
return x
class DFLV3Rewrite(DFPatternCallback):
def __init__(self):
super().__init__()
self.x = wildcard()
self.reshape = is_op("reshape")(self.x)
self.softmax = is_op("nn.softmax")(self.reshape).has_attr({"axis": 1})
self.conv_weight = is_constant()
self.conv = is_op("nn.conv2d")(self.softmax, self.conv_weight)
self.reshape2 = is_op("reshape")(self.conv)
self.pattern = self.reshape2
def callback(self, pre, post, node_map):
x = node_map[self.x][0]
conv_weight = node_map[self.conv_weight][0]
b, c, a = _transform.InferTypeLocal(x).shape # batch, channels, anchors
conv_weight_shape = _transform.InferTypeLocal(conv_weight).shape
assert conv_weight_shape[0] == conv_weight_shape[2] == conv_weight_shape[3] == 1
x = yolo_dfl(x, conv_weight_shape[1], (b, c, a), version="v3")
return x
mod["main"] = rewrite(DFLV1Rewrite(), mod["main"])
mod["main"] = rewrite(DFLV2Rewrite(), mod["main"])
mod["main"] = rewrite(DFLV3Rewrite(), mod["main"])
mod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), float32] span=/conv0/Conv.conv0.weight:0:0 */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1]) /* ty=Tensor[(1, 16, 48, 80), float32] span=/conv0/Conv:0:0 */;
%1 = nn.conv2d(%0, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv1/Conv.conv1.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv1/Conv:0:0 */;
%2 = reshape(%1, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape:0:0 */;
%3 = nn.conv2d(%0, meta[relay.Constant][2] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv2/Conv.conv2.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv2/Conv:0:0 */;
%4 = reshape(%3, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape_1:0:0 */;
%5 = nn.conv2d(%0, meta[relay.Constant][3] /* ty=Tensor[(100, 16, 1, 1), float32] span=/conv3/Conv.conv3.weight:0:0 */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1]) /* ty=Tensor[(1, 100, 48, 80), float32] span=/conv3/Conv:0:0 */;
%6 = reshape(%5, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 6000), float32] span=/Reshape_2:0:0 */;
%7 = vta_special.yolo_dfl(%2, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v1", "channel"=16});
%8 = vta_special.yolo_dfl(%4, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v2", "channel"=16});
%9 = vta_special.yolo_dfl(%6, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 6000], "version"="v3", "channel"=16});
(%7, %8, %9) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
origin_mod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), float32] span=/conv0/Conv.conv0.weight:0:0 */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1]) /* ty=Tensor[(1, 16, 48, 80), float32] span=/conv0/Conv:0:0 */;
%1 = nn.conv2d(%0, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv1/Conv.conv1.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv1/Conv:0:0 */;
%2 = reshape(%1, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape:0:0 */;
%3 = reshape(%2, newshape=[1, 4, 16, 3840]) /* ty=Tensor[(1, 4, 16, 3840), float32] span=/v1/Reshape:0:0 */;
%4 = transpose(%3, axes=[0, 3, 1, 2]) /* ty=Tensor[(1, 3840, 4, 16), float32] span=/v1/Transpose:0:0 */;
%5 = nn.softmax(%4, axis=3) /* ty=Tensor[(1, 3840, 4, 16), float32] span=/v1/Softmax:0:0 */;
%6 = transpose(%5, axes=[0, 3, 2, 1]) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v1/Transpose_1:0:0 */;
%7 = nn.conv2d(%6, meta[relay.Constant][2] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v1/conv/Conv.v1.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 3840), float32] span=/v1/conv/Conv:0:0 */;
%8 = nn.conv2d(%0, meta[relay.Constant][3] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv2/Conv.conv2.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv2/Conv:0:0 */;
%9 = reshape(%8, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape_1:0:0 */;
%10 = reshape(%9, newshape=[1, 4, 16, 3840]) /* ty=Tensor[(1, 4, 16, 3840), float32] span=/v2/Reshape:0:0 */;
%11 = transpose(%10, axes=[0, 2, 1, 3]) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v2/Transpose:0:0 */;
%12 = nn.softmax(%11, axis=1) /* ty=Tensor[(1, 16, 4, 3840), float32] span=/v2/Softmax:0:0 */;
%13 = nn.conv2d(%12, meta[relay.Constant][4] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v2/conv/Conv.v2.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 3840), float32] span=/v2/conv/Conv:0:0 */;
%14 = nn.conv2d(%0, meta[relay.Constant][5] /* ty=Tensor[(100, 16, 1, 1), float32] span=/conv3/Conv.conv3.weight:0:0 */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1]) /* ty=Tensor[(1, 100, 48, 80), float32] span=/conv3/Conv:0:0 */;
%15 = reshape(%14, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 6000), float32] span=/Reshape_2:0:0 */;
%16 = reshape(%15, newshape=[1, 16, 4, 6000]) /* ty=Tensor[(1, 16, 4, 6000), float32] span=/v3/Reshape:0:0 */;
%17 = nn.softmax(%16, axis=1) /* ty=Tensor[(1, 16, 4, 6000), float32] span=/v3/Softmax:0:0 */;
%18 = nn.conv2d(%17, meta[relay.Constant][6] /* ty=Tensor[(1, 16, 1, 1), float32] span=/v3/conv/Conv.v3.conv.weight:0:0 */, padding=[0, 0, 0, 0], channels=1, kernel_size=[1, 1]) /* ty=Tensor[(1, 1, 4, 6000), float32] span=/v3/conv/Conv:0:0 */;
%19 = reshape(%7, newshape=[1, 4, 3840]) /* ty=Tensor[(1, 4, 3840), float32] span=/v1/Reshape_1:0:0 */;
%20 = reshape(%13, newshape=[1, 4, 3840]) /* ty=Tensor[(1, 4, 3840), float32] span=/v2/Reshape_1:0:0 */;
%21 = reshape(%18, newshape=[1, 4, 6000]) /* ty=Tensor[(1, 4, 6000), float32] span=/v3/Reshape_1:0:0 */;
(%19, %20, %21) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
data = np.random.normal(0, 1, size=shape).astype("float32")
with torch.no_grad():
torch_outputs = [o.numpy() for o in model(torch.from_numpy(data))]
target = 'llvm'
dev = tvm.device(target, 0)
# 原始模型
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(origin_mod, target, params=params)
func = lib[lib.libmod_name]
module = tvm.contrib.graph_executor.GraphModule(func(dev))
module.run(**{input_name: data})
outputs1 = [module.get_output(k).numpy() for k in range(3)]
# 重写后的模型
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target, params=params)
func = lib[lib.libmod_name]
module = tvm.contrib.graph_executor.GraphModule(func(dev))
module.run(**{input_name: data})
outputs2 = [module.get_output(k).numpy() for k in range(3)]
[np.testing.assert_allclose(torch_outputs[k], outputs1[k], rtol=1e-07, atol=1e-5) for k in range(3)]
[np.testing.assert_allclose(torch_outputs[k], outputs2[k], rtol=1e-07, atol=1e-5) for k in range(3)]
[np.testing.assert_equal(outputs1[k], outputs2[k]) for k in range(3)];
One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
mod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), float32] span=/conv0/Conv.conv0.weight:0:0 */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1]) /* ty=Tensor[(1, 16, 48, 80), float32] span=/conv0/Conv:0:0 */;
%1 = nn.conv2d(%0, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv1/Conv.conv1.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv1/Conv:0:0 */;
%2 = reshape(%1, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape:0:0 */;
%3 = nn.conv2d(%0, meta[relay.Constant][2] /* ty=Tensor[(64, 16, 1, 1), float32] span=/conv2/Conv.conv2.weight:0:0 */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 48, 80), float32] span=/conv2/Conv:0:0 */;
%4 = reshape(%3, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), float32] span=/Reshape_1:0:0 */;
%5 = nn.conv2d(%0, meta[relay.Constant][3] /* ty=Tensor[(100, 16, 1, 1), float32] span=/conv3/Conv.conv3.weight:0:0 */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1]) /* ty=Tensor[(1, 100, 48, 80), float32] span=/conv3/Conv:0:0 */;
%6 = reshape(%5, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 6000), float32] span=/Reshape_2:0:0 */;
%7 = vta_special.yolo_dfl(%2, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v1", "channel"=16});
%8 = vta_special.yolo_dfl(%4, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v2", "channel"=16});
%9 = vta_special.yolo_dfl(%6, __dict__={"x_scale"=-1, "x_split"=-1, "in_shape"=[1, 64, 6000], "version"="v3", "channel"=16});
(%7, %8, %9) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
from PIL import Image
import numpy as np
image = np.random.normal(0, 1, size=(48, 80, 3)).astype("uint8")
mean = (128,)
std = (256,)
data = (image - mean)/std
data = data.transpose((2, 0, 1))
data = np.expand_dims(data, 0).astype("float32")
images = np.expand_dims(image, 0)
images.tofile(f"{root_dir}/input.bin")
Image.fromarray(image).resize((112, 112))
from dataclasses import dataclass
@dataclass
class Dataset:
input_name: str
shape: tuple
def __iter__(self):
for _ in range(2):
yield {self.input_name: data}
# for _ in range(50):
# yield {self.input_name: np.random.normal(0, 1, size=self.shape).astype("float32")}
dataset = Dataset(input_name, shape)
with relay.quantize.qconfig(
calibrate_mode="kl_divergence",
weight_scale="max",
# vta_adjust_scale = True, # 跑 vta 时必须开启
# vta_relu6_flag = True,
# weight_per_channel_quantization=True, # 逐通道量化
# prelu_fuse=True,
skip_conv_layers=[],
skip_dense_layer=False,):
# qmod, record_graphs = relay.quantize.quantize_debug(mod, params, dataset)
qmod = relay.quantize.quantize(mod, params, dataset)
# QNN 量化的模型
origin_qmod = deepcopy(qmod)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(origin_qmod, target, params=params)
func = lib[lib.libmod_name]
module = tvm.contrib.graph_executor.GraphModule(func(dev))
module.run(**{input_name: data})
origin_qoutputs = [module.get_output(k).numpy() for k in range(3)]
origin_qmod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = multiply(%data, 257.869f /* ty=float32 */) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%1 = round(%0) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%2 = clip(%1, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%3 = cast(%2, dtype="int8") /* ty=Tensor[(1, 3, 48, 80), int8] */;
%4 = nn.conv2d(%3, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 16, 48, 80), int32] */;
%5 = cast(%4, dtype="int64") /* ty=Tensor[(1, 16, 48, 80), int64] */;
%6 = fixed_point_multiply(%5, multiplier=1712233856, shift=-8) /* ty=Tensor[(1, 16, 48, 80), int64] */;
%7 = clip(%6, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 16, 48, 80), int64] */;
%8 = cast(%7, dtype="int32") /* ty=Tensor[(1, 16, 48, 80), int32] */;
%9 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%10 = annotation.stop_fusion(%9) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%11 = nn.conv2d(%10, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%12 = cast(%11, dtype="int64") /* ty=Tensor[(1, 64, 48, 80), int64] */;
%13 = fixed_point_multiply(%12, multiplier=1878688768, shift=-8) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%14 = clip(%13, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%15 = cast(%14, dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%16 = cast(%15, dtype="int8") /* ty=Tensor[(1, 64, 48, 80), int8] */;
%17 = annotation.stop_fusion(%16) /* ty=Tensor[(1, 64, 48, 80), int8] */;
%18 = reshape(%17, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), int8] */;
%19 = cast(%18, dtype="float32") /* ty=Tensor[(1, 64, 3840), float32] */;
%20 = multiply(%19, 0.00317437f /* ty=float32 */) /* ty=Tensor[(1, 64, 3840), float32] */;
%21 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%22 = annotation.stop_fusion(%21) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%23 = nn.conv2d(%22, meta[relay.Constant][2] /* ty=Tensor[(64, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%24 = cast(%23, dtype="int64") /* ty=Tensor[(1, 64, 48, 80), int64] */;
%25 = fixed_point_multiply(%24, multiplier=1496169600, shift=-8) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%26 = clip(%25, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%27 = cast(%26, dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%28 = cast(%27, dtype="int8") /* ty=Tensor[(1, 64, 48, 80), int8] */;
%29 = annotation.stop_fusion(%28) /* ty=Tensor[(1, 64, 48, 80), int8] */;
%30 = reshape(%29, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 3840), int8] */;
%31 = cast(%30, dtype="float32") /* ty=Tensor[(1, 64, 3840), float32] */;
%32 = multiply(%31, 0.00398762f /* ty=float32 */) /* ty=Tensor[(1, 64, 3840), float32] */;
%33 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%34 = annotation.stop_fusion(%33) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%35 = nn.conv2d(%34, meta[relay.Constant][3] /* ty=Tensor[(100, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 100, 48, 80), int32] */;
%36 = cast(%35, dtype="int64") /* ty=Tensor[(1, 100, 48, 80), int64] */;
%37 = fixed_point_multiply(%36, multiplier=1540201216, shift=-8) /* ty=Tensor[(1, 100, 48, 80), int64] */;
%38 = clip(%37, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 100, 48, 80), int64] */;
%39 = cast(%38, dtype="int32") /* ty=Tensor[(1, 100, 48, 80), int32] */;
%40 = cast(%39, dtype="int8") /* ty=Tensor[(1, 100, 48, 80), int8] */;
%41 = annotation.stop_fusion(%40) /* ty=Tensor[(1, 100, 48, 80), int8] */;
%42 = reshape(%41, newshape=[1, 64, -1]) /* ty=Tensor[(1, 64, 6000), int8] */;
%43 = cast(%42, dtype="float32") /* ty=Tensor[(1, 64, 6000), float32] */;
%44 = multiply(%43, 0.00387525f /* ty=float32 */) /* ty=Tensor[(1, 64, 6000), float32] */;
%45 = vta_special.yolo_dfl(%20, __dict__={"x_scale"=-1, "in_shape"=[1, 64, 3840], "version"="v1", "x_split"=-1, "channel"=16}) /* ty=Tensor[(1, 4, 3840), float32] */;
%46 = vta_special.yolo_dfl(%32, __dict__={"x_scale"=-1, "in_shape"=[1, 64, 3840], "version"="v2", "x_split"=-1, "channel"=16}) /* ty=Tensor[(1, 4, 3840), float32] */;
%47 = vta_special.yolo_dfl(%44, __dict__={"x_scale"=-1, "in_shape"=[1, 64, 6000], "version"="v3", "x_split"=-1, "channel"=16}) /* ty=Tensor[(1, 4, 6000), float32] */;
(%45, %46, %47) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
class YoloDFLPredictRewrite(DFPatternCallback):
"""融合 (reshape)+cast+multiply+vta_special.vta_yolo_dfl_predict 以更新 x_scale
"""
def __init__(self):
super().__init__()
self.x = wildcard()
self.reshape = is_op("reshape")(self.x)
self.cast = is_op("cast")(self.reshape|self.x).has_attr({"dtype": "float32"})
self.data_scale = is_constant()
self.multiply = is_op("multiply")(self.cast, self.data_scale)
self.dfl_predict_call = is_op("vta_special.yolo_dfl")(self.multiply)
self.pattern = self.dfl_predict_call
def callback(self, pre, post, node_map):
x = node_map[self.x][0]
data_scale = node_map[self.data_scale][0]
dfl_predict_call = node_map[self.dfl_predict_call][0]
return yolo_dfl(
x, dfl_predict_call.attrs.channel,
dfl_predict_call.attrs.in_shape,
dfl_predict_call.attrs.version,
float(data_scale.data.numpy())
)
qmod["main"] = rewrite(YoloDFLPredictRewrite(), qmod["main"])
qmod.show()
def @main(%data: Tensor[(1, 3, 48, 80), float32] /* ty=Tensor[(1, 3, 48, 80), float32] span=/conv0/Conv.data:0:0 */) -> (Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) {
%0 = multiply(%data, 257.869f /* ty=float32 */) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%1 = round(%0) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%2 = clip(%1, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 3, 48, 80), float32] */;
%3 = cast(%2, dtype="int8") /* ty=Tensor[(1, 3, 48, 80), int8] */;
%4 = nn.conv2d(%3, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=16, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 16, 48, 80), int32] */;
%5 = cast(%4, dtype="int64") /* ty=Tensor[(1, 16, 48, 80), int64] */;
%6 = fixed_point_multiply(%5, multiplier=1712233856, shift=-8) /* ty=Tensor[(1, 16, 48, 80), int64] */;
%7 = clip(%6, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 16, 48, 80), int64] */;
%8 = cast(%7, dtype="int32") /* ty=Tensor[(1, 16, 48, 80), int32] */;
%9 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%10 = annotation.stop_fusion(%9) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%11 = nn.conv2d(%10, meta[relay.Constant][1] /* ty=Tensor[(64, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%12 = cast(%11, dtype="int64") /* ty=Tensor[(1, 64, 48, 80), int64] */;
%13 = fixed_point_multiply(%12, multiplier=1878688768, shift=-8) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%14 = clip(%13, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%15 = cast(%14, dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%16 = cast(%15, dtype="int8") /* ty=Tensor[(1, 64, 48, 80), int8] */;
%17 = annotation.stop_fusion(%16) /* ty=Tensor[(1, 64, 48, 80), int8] */;
%18 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%19 = annotation.stop_fusion(%18) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%20 = nn.conv2d(%19, meta[relay.Constant][2] /* ty=Tensor[(64, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%21 = cast(%20, dtype="int64") /* ty=Tensor[(1, 64, 48, 80), int64] */;
%22 = fixed_point_multiply(%21, multiplier=1496169600, shift=-8) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%23 = clip(%22, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 64, 48, 80), int64] */;
%24 = cast(%23, dtype="int32") /* ty=Tensor[(1, 64, 48, 80), int32] */;
%25 = cast(%24, dtype="int8") /* ty=Tensor[(1, 64, 48, 80), int8] */;
%26 = annotation.stop_fusion(%25) /* ty=Tensor[(1, 64, 48, 80), int8] */;
%27 = cast(%8, dtype="int8") /* ty=Tensor[(1, 16, 48, 80), int8] */;
%28 = annotation.stop_fusion(%27) /* ty=Tensor[(1, 16, 48, 80), int8] */;
%29 = nn.conv2d(%28, meta[relay.Constant][3] /* ty=Tensor[(100, 16, 1, 1), int8] */, padding=[0, 0, 0, 0], channels=100, kernel_size=[1, 1], out_dtype="int32") /* ty=Tensor[(1, 100, 48, 80), int32] */;
%30 = cast(%29, dtype="int64") /* ty=Tensor[(1, 100, 48, 80), int64] */;
%31 = fixed_point_multiply(%30, multiplier=1540201216, shift=-8) /* ty=Tensor[(1, 100, 48, 80), int64] */;
%32 = clip(%31, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 100, 48, 80), int64] */;
%33 = cast(%32, dtype="int32") /* ty=Tensor[(1, 100, 48, 80), int32] */;
%34 = cast(%33, dtype="int8") /* ty=Tensor[(1, 100, 48, 80), int8] */;
%35 = annotation.stop_fusion(%34) /* ty=Tensor[(1, 100, 48, 80), int8] */;
%36 = vta_special.yolo_dfl(%17, __dict__={"x_scale"=0.00317437f, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v1", "channel"=16});
%37 = vta_special.yolo_dfl(%26, __dict__={"x_scale"=0.00398762f, "x_split"=-1, "in_shape"=[1, 64, 3840], "version"="v2", "channel"=16});
%38 = vta_special.yolo_dfl(%35, __dict__={"x_scale"=0.00387525f, "x_split"=-1, "in_shape"=[1, 64, 6000], "version"="v3", "channel"=16});
(%36, %37, %38) /* ty=(Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 3840), float32], Tensor[(1, 4, 6000), float32]) */
}
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(qmod, target, params=params)
func = lib[lib.libmod_name]
module = tvm.contrib.graph_executor.GraphModule(func(dev))
module.run(**{input_name: data})
qoutputs = [module.get_output(k).numpy() for k in range(3)]
[np.testing.assert_almost_equal(origin_qoutputs[k], qoutputs[k]) for k in range(3)];