Nerotos opened a new issue, #14835:
URL: https://github.com/apache/tvm/issues/14835
Hello,
I have started working on QuantizeInferCorrectLayout as I need to use
Quantize with an NHWC layout, but I can't figure out how to implement this. I
started by just copying over the version requantize.cc
### Expected behavior
After the ConvertLayoutpass I want to have all data in NHWC format
throughout the IR.
### Actual behavior
Some elements are not in the correct layout. See calculation of %5 and %6.
```
fn (%layer1_input: Tensor[(1, 16, 16, 16), float32] /* ty=Tensor[(1, 16, 16,
16), float32] */, %c1_weight: Tensor[(16, 16, 3, 3), float32] /* ty=Tensor[(16,
16, 3, 3), float32] */, %c1_bias: Tensor[(16), float32] /* ty=Tensor[(16),
float32] */) -> Tensor[(1, 16, 14, 14), uint8] {
%0 = qnn.quantize(%layer1_input, 0.00786161f /* ty=float32 */, 0 /*
ty=int32 */, out_dtype="uint8", axis=1) /* ty=Tensor[(1, 16, 16, 16), uint8] */;
%1 = qnn.quantize(%c1_weight, meta[relay.Constant][0] /* ty=Tensor[(16),
float32] */, 0 /* ty=int32 */, out_dtype="int8", axis=0) /* ty=Tensor[(16, 16,
3, 3), int8] */;
%2 = layout_transform(%0, src_layout="NCHW", dst_layout="NHWC") /*
ty=Tensor[(1, 16, 16, 16), uint8] */;
%3 = layout_transform(%1, src_layout="OIHW", dst_layout="HWIO") /*
ty=Tensor[(3, 3, 16, 16), int8] */;
%4 = qnn.conv2d(%2, %3, 0 /* ty=int32 */, 0 /* ty=int32 */, 0.00786161f /*
ty=float32 */, meta[relay.Constant][0] /* ty=Tensor[(16), float32] */,
padding=[0, 0, 0, 0], channels=16, kernel_size=[3, 3], data_layout="NHWC",
kernel_layout="HWIO", out_dtype="int32") /* ty=Tensor[(1, 14, 14, 16), int32]
*/;
%5 = layout_transform(%4, src_layout="NHWC", dst_layout="NCHW") /*
ty=Tensor[(1, 16, 14, 14), int32] */;
%6 = qnn.quantize(%c1_bias, meta[relay.Constant][1] /* ty=Tensor[(16),
float32] */, 0 /* ty=int32 */, out_dtype="int32", axis=0) /* ty=Tensor[(16),
int32] */;
%7 = nn.bias_add(%5, %6) /* ty=Tensor[(1, 16, 14, 14), int32] */;
%8 = qnn.requantize(%7, meta[relay.Constant][2] /* ty=Tensor[(16),
float32] */, 0 /* ty=int32 */, 0.0175786f /* ty=float32 */, 64 /* ty=int32 */,
axis=1, out_dtype="int32") /* ty=Tensor[(1, 16, 14, 14), int32] */;
%9 = clip(%8, a_min=0f, a_max=255f) /* ty=Tensor[(1, 16, 14, 14), int32]
*/;
cast(%9, dtype="uint8") /* ty=Tensor[(1, 16, 14, 14), uint8] */
} /* ty=fn (Tensor[(1, 16, 16, 16), float32], Tensor[(16, 16, 3, 3),
float32], Tensor[(16), float32]) -> Tensor[(1, 16, 14, 14), uint8] */
```
### Environment
Linux OS
TVM main branch
### Steps to reproduce
Current code for QuantizeInferCorrectLayout, from requantize. Copy into
tmv/src/relay/qnn/op/quantize.cc:
```
InferCorrectLayoutOutput QuantizeInferCorrectLayout(const Attrs& attrs,
const Array<Layout>&
new_in_layouts,
const Array<Layout>&
old_in_layouts,
const
Array<tvm::relay::Type>& old_in_types) {
const auto* attrs_ptr = attrs.as<QuantizeAttrs>();
ICHECK(attrs_ptr);
ObjectPtr<QuantizeAttrs> param = make_object<QuantizeAttrs>(*attrs_ptr);
Array<Array<IndexExpr>> old_in_shapes;
for (auto old_in_t : old_in_types) {
ICHECK(old_in_t.as<TensorTypeNode>());
old_in_shapes.push_back(old_in_t.as<TensorTypeNode>()->shape);
}
Array<Layout> input_layouts, output_layouts;
if (new_in_layouts.defined()) {
// Adapt to new layout. The axis has to change.
// Record original reduce axis. Convert to the modified layout axis.
//ICHECK_EQ(new_in_layouts.size(), 3);
ICHECK_EQ(old_in_layouts.size(), 3);
// 1) Get the axis.
int axis = param->axis;
axis = (axis == -1) ? old_in_shapes[0].size() - 1 : axis;
// 2) Collect the original axis
std::string old_dim = old_in_layouts[0][axis].name();
// 3) Collect the new axes by walking new_layout.
tvm::Integer new_axis;
std::string new_layout_string = "";
int axis_index = 0;
for (auto iter_var : new_in_layouts[0]->axes) {
const auto& layout_axis = LayoutAxis::Get(iter_var);
const std::string& layout_dim = layout_axis.name();
if (old_dim == layout_dim) {
new_axis = tvm::Integer(axis_index);
}
if (layout_axis.IsPrimal()) {
new_layout_string += layout_dim;
axis_index++;
} else {
// Propagate layout if input_zero_point and input_scale are scalar
values.
ICHECK_GE(old_in_types.size(), 3);
if (IsScalarType(old_in_types[1]) && IsScalarType(old_in_types[2])) {
new_layout_string +=
std::to_string(new_in_layouts[0].FactorOf(layout_axis)) + layout_dim;
axis_index++;
}
}
}
// 4) Set the new axis and layout.
Layout new_layout = Layout(new_layout_string);
// Fill the layouts of remaining input tensors - scales and zero points.
The layouts of these
// tensors can be treated as channel layout.
Layout channel_layout = Layout("C");
input_layouts = {new_layout, channel_layout, channel_layout,
channel_layout, channel_layout};
output_layouts = {new_layout};
param->axis = new_axis.IntValue();
} else if (old_in_layouts.defined()) {
// If the new layout is undefined, set the old layout as the inferred
layout.
ICHECK_EQ(old_in_layouts.size(), 3);
Layout old_layout = old_in_layouts[0];
// Fill the layouts of remaining input tensors - scales and zero points.
The layouts of these
// tensors can be treated as channel layout.
Layout channel_layout = Layout("C");
input_layouts = {old_layout, channel_layout, channel_layout,
channel_layout, channel_layout};
output_layouts = {old_layout};
} else {
// Set the layouts to undef.
Layout undef = Layout::Undef();
input_layouts = Array<Layout>(3, undef);
output_layouts = {undef};
}
return InferCorrectLayoutOutput(input_layouts, output_layouts,
Attrs(param));
}
```
An example script to import a model from pytorch and apply the
transformations:
```
import numpy as np
from torch import nn
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization.observer import MinMaxObserver
import torch
import os
from tvm import relay
import tvm
batch_size = 1
input_channels = 16
input_height = 16
input_width = 16
output_channels = 16
kernel_size = 3
stride = 1
num_classes = 1
bias = False
class ConvModel(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size,
num_classes=0):
super().__init__()
self.quant = torch.ao.quantization.QuantStub()
self.c1 = nn.Conv2d(in_channels, out_channels, kernel_size)
def forward(self, x):
x = self.quant(x)
x = self.c1(x)
return x
# Quantize the model
input_fp = torch.rand(1, input_channels, input_height, input_width)
model = ConvModel(input_channels, output_channels, kernel_size)
model.eval()
# Specify quantization configuration
# Start with simple min/max range estimation and per-tensor quantization of
weights
#model.qconfig = torch.ao.quantization.get_default_qconfig("x86")
model.qconfig = torch.ao.quantization.QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8),
weight=MinMaxObserver.with_args(dtype=torch.qint8))
torch.ao.quantization.prepare(model, inplace=True)
# Pseudo-Calibration
model(input_fp)
#Convert to quantized model
torch.ao.quantization.convert(model, inplace=True)
# Save the model.
torch.jit.save(torch.jit.script(model),
"conv2d_model_scripted_quantized.pth")
os.system("rm -rf model.tar dev/ include/ generated-project/")
os.system("mkdir -p include")
#Generate expected output data
input_matrix = np.random.randint(0, high=127, size=(1, input_channels,
input_height, input_width), dtype=np.int8)
input_float = torch.rand(1, input_channels, input_height, input_width)
input_quant = input_float
expected_output = model(input_quant)
##################################
# Compiling the model with TVM
# --------------------------------
#
# For importing from Torch we need to trace the model beforehand
input_name = "layer1_input"
input_shapes = [(input_name, (batch_size, input_channels, input_height,
input_width))]
input_dtype = "uint8"
script_module = torch.jit.trace(model, input_quant).eval()
mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
mod = gemmini.preprocess_pass(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.ConvertLayout({"qnn.conv2d": ["NHWC", "HWIO"],
"nn.conv2d": ["NHWC", "HWIO"]})(mod)
mod = relay.transform.SimplifyExpr()(mod)
```
### Triage
* needs-triage
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
