【AI系统】TVM 实践案例

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import data_set

# 模型加载
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = torchvision.models.resnet18(pretrained= False)
names = [item[0] for item in net._modules.items()]
model_weight_path = "./weights/resnet18-f37072fd.pth"
net.load_state_dict(torch.load(model_weight_path), strict=False)

# 指定融合层
fused_layer = [['model_fp32.conv1', 'model_fp32.bn1', 'model_fp32.relu'],
    ['model_fp32.layer1.0.conv1', 'model_fp32.layer1.0.bn1', 'model_fp32.layer1.0.relu'],
    ['model_fp32.layer1.0.conv2', 'model_fp32.layer1.0.bn2'],
    ['model_fp32.layer1.1.conv1', 'model_fp32.layer1.1.bn1', 'model_fp32.layer1.1.relu'],
    ['model_fp32.layer1.1.conv2', 'model_fp32.layer1.1.bn2'],
    ['model_fp32.layer2.0.conv1', 'model_fp32.layer2.0.bn1', 'model_fp32.layer2.0.relu'],
    ['model_fp32.layer2.0.conv2', 'model_fp32.layer2.0.bn2'],
    ['model_fp32.layer2.1.conv1', 'model_fp32.layer2.1.bn1', 'model_fp32.layer2.1.relu'],
    ['model_fp32.layer2.1.conv2', 'model_fp32.layer2.1.bn2'],
    ['model_fp32.layer3.0.conv1', 'model_fp32.layer3.0.bn1', 'model_fp32.layer3.0.relu'],
    ['model_fp32.layer3.0.conv2', 'model_fp32.layer3.0.bn2'],
    ['model_fp32.layer3.1.conv1', 'model_fp32.layer3.1.bn1', 'model_fp32.layer3.1.relu'],
    ['model_fp32.layer3.1.conv2', 'model_fp32.layer3.1.bn2'],
    ['model_fp32.layer4.0.conv1', 'model_fp32.layer4.0.bn1', 'model_fp32.layer4.0.relu'],
    ['model_fp32.layer4.0.conv2', 'model_fp32.layer4.0.bn2'],
    ['model_fp32.layer4.1.conv1', 'model_fp32.layer4.1.bn1', 'model_fp32.layer4.1.relu'],
    ['model_fp32.layer4.1.conv2', 'model_fp32.layer4.1.bn2'],
    ]

# 插入量化与反量化节点
class QuantizedResNet18(nn.Module):
    def __init__(self, model_fp32):
        super(QuantizedResNet18, self).__init__()
        # QuantStub converts tensors from floating point to quantized.
        # This will only be used for inputs.
        self.quant = torch.quantization.QuantStub()
        # DeQuantStub converts tensors from quantized to floating point.
        # This will only be used for outputs.
        self.dequant = torch.quantization.DeQuantStub()
        # FP32 model
        self.model_fp32 = model_fp32

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        x = self.quant(x)
        x = self.model_fp32.forward(x)
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        return x


quantized_model = QuantizedResNet18(net)
quantized_model.eval()

# 层融合
quantized_model = torch.ao.quantization.fuse_modules(quantized_model, fused_layer)
quantized_model.train()

# 量化配置
quantization_config = torch.quantization.QConfig(activation=torch.quantization.MinMaxObserver.with_args(dtype=torch.quint16), weight=torch.quantization.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_affine))
quantized_model.qconfig = quantization_config
quantized_model=torch.ao.quantization.prepare_qat(quantized_model, inplace=True)
quantized_model.to(device)

#超参数
batch_size = 64
lr = 0.0001
epoch = 10


trainloader,testloader,test_num = data_set.load_data(batch_size)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(quantized_model.parameters(), lr=lr)


# 训练
best_acc = 0.0
save_path_fp32 = './weights/resNet18_fp32.pth'
for epoch in range(epoch):
    # training
    quantized_model.train()
    running_loss = 0.0
    for step, data in enumerate(trainloader, start=0):
        images, labels = data
        optimizer.zero_grad()
        logits = quantized_model(images.to(device))
        loss = loss_function(logits, labels.to(device))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        rate = (step + 1) / len(trainloader)
        a = "*" * int(rate * 50)
        b = "." * int((1 - rate) * 50)
        print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate * 100), a, b, loss), end="")

    quantized_model.eval()  # change into test model
    # int_model = torch.ao.quantization.convert(quantized_model, inplace=True)
    # int_model.eval()
    # 以上代码可以将模型转换为 int8 模型进行测试
    acc = 0.0  # accumulate accurate number / epoch
    with torch.no_grad():
        for val_data in testloader:
            val_images, val_labels = val_data
            outputs = quantized_model(val_images.to(device))  # eval model only have last output layer
            # loss = loss_function(outputs, test_labels)
            predict_y = torch.max(outputs, dim=1)[1]
            acc += (predict_y == val_labels.to(device)).sum().item()
        val_accurate = acc / test_num
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(quantized_model.state_dict(), save_path_fp32)
        print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f' %
            (epoch + 1, running_loss / step, val_accurate))

save_path_int8 = './weights/resNet18_INT16.pth'
quantized_model.eval()  # change into test model

# 转换为 int8 模型进行保存
int_model = torch.quantization.convert(quantized_model, inplace=True)
torch.save(int_model.state_dict(), save_path_int8)

from torch.quantization.quantize_fx import prepare_qat_fx,convert_fx
import torch.quantization.observer as observer

#将模型转换为 QAT 版本模型，其他训练过程与正常训练类似
def qat_version_model(model):
    qconfig_dict = {
        # Global Config
        "":torch.ao.quantization.get_default_qat_qconfig('qnnpack'), #全局量化配置

        # # Disable by layer-name
        # "module_name": [(m, None) for m in disable_layers],

        # Or disable by layer-type
        "object_type": [
            (PositionalEmbedding, None),   #想要跳过量化层，可以设置为 None
            (torch.nn.Softmax, softmax_qconfig), #指定与全局量化配置不同的量化配置
            ......
        ],
    }
    model_to_quantize = copy.deepcopy(model)
    model_fp32_prepared = prepare_qat_fx(
        model_to_quantize, qconfig_dict)
    return model_fp32_prepared

qat_model=qat_version_model(float_model)

#********模型训练**********
for i in range(epoch_nums):
    train(qat_model,train_data_loader)
    ......

#********模型验证**********
quantized_model = convert_fx(qat_model)
evaluate(quantized_model,eval_data_loader)
......

# 保存量化版本的模型

# 以 torchScript 的形式
trace_model = torch.jit.trace(quantized_model, (input1, input2, input3))
torch.jit.save(trace_model, 'trace_model.pt')

#以 onnx 的形式
torch.onnx.export(quantized_model, (input1, input2, input3), "quantized_model.onnx",
input_names=['input1', 'input2', 'input3'], output_names=['result'],
opset_version=16, export_params=True)

xxxfunc(): argument 'xxx' (position 1) must be xxx, not Proxy

@torch.fx.wrap
def abs_fx(dim1:int,dim2:int):
    return abs(dim1 - dim2)
    
@torch.fx.wrap
def ones_fx(dim1:int,dim2:int):
    return torch.ones(dim1, dim2)
    
@torch.fx.wrap
def triu_fx(dim1:int,dim2:int):
    return torch.triu(dim1,dim2)

def _quantized_conv2d(with_relu=False):
    def _impl(inputs, _):
        # refer to src/ATen/native/quantized/cpu/qconv.cpp
        # inputs[0]: input tensor
        # inputs[1]: (weight, scale, zero_point, bias)
        # inputs[2-5]: stride, padding, dilation, groups
        # inputs[6]: output_scale
        # inputs[7]: output_zero_point
        # inputs[8]: input_scale (added manually by frontend)   input 的 scale 和 zero_point 在 QTensor 的属性中，需要提前运行_get_quant_param_for_input()函数获得。该函数有一个 output_quant_param_indices 的 dict，记录了算子输出结果的 scale 和 zero_point，作为下一个算子 input 的 scale 和 zero_point
        # inputs[9]: input_zero_point (added manually by frontend)
        conv_params = inputs[1]
        weight = conv_params[0]
        weight_scale = conv_params[1]
        weight_zero_point = conv_params[2]
        bias = conv_params[3]

        if len(conv_params) > 4:
            # Torch 1.6 or newer case
            strides = conv_params[4]
            padding = conv_params[5]
            dilation = conv_params[6]
            groups = conv_params[7]

            output_scale = _expr.const(inputs[2])
            output_zero_point = _expr.const(inputs[3])

            assert len(inputs) == 6, "Input quant params not found in op inputs"

            # These are manually added by add_input_quant_params_to_op_inputs above
            # In torch, they are retrieved from QTensor data structure at runtime
            input_scale = _expr.const(inputs[4])
            input_zero_point = _expr.const(inputs[5])
        else:
            strides = inputs[2]
            padding = inputs[3]
            dilation = inputs[4]
            groups = inputs[5]
            output_scale = _expr.const(inputs[6])
            output_zero_point = _expr.const(inputs[7])

            assert len(
                inputs) == 10, "Input quant params not found in op inputs"

            input_scale = _expr.const(inputs[8])
            input_zero_point = _expr.const(inputs[9])

        weight_shape = infer_shape(weight)
        kernel_size = (weight_shape[2], weight_shape[3])
        out_channels = weight_shape[0]

        if padding[0] != 0 or padding[1] != 0:
            pad_val = _get_scalar(input_zero_point)
            inp = _op.nn.pad(
                inputs[0],
                pad_width=((0, 0), (0, 0),
                           (padding[0], padding[0]), (padding[1], padding[1])),
                pad_value=float(pad_val), #padding 操作的值是输入的 zero_point，而不直接是 0
            )
        else:
            inp = inputs[0]

        # padding is (0, 0) because we did explicit pad op with
        # pad value being zero point above
        conv_out = relay.qnn.op.conv2d(  #对接了 qnn 的 conv2d 算子，该算子的实现在 cpp 端，并没有直接实现量化 conv2d 的 kernel，而是转换为一系列 relay 操作
            inp,
            weight,
            input_zero_point,
            weight_zero_point,
            input_scale,
            weight_scale,
            kernel_size=kernel_size,
            dilation=dilation,
            strides=strides,
            padding=(0, 0),
            groups=groups,
            channels=out_channels,
        )

        #在计算完 quantized::conv2d 后，需要进行再量化
        return _do_bias_and_requantize(
            conv_out, bias, input_scale, weight_scale, output_scale, output_zero_point, with_relu
        )

    return _impl

A = scale_a x (QA - zp_A)

C(m, n) = Sigma(k) (A(m, k) * W(n, k)) // 表示在 k axis 进行 reduce

Sigma(k) ([QA(m, k) - zp_a] * [QW(n, k) - zp_w])

Sigma(k) QA(m, k) * QW(n, k)                         // Term1
- Sigma(k) zp_w * QA(m, k)                             // Term2
- Sigma(k) zp_a * QW(n, k)                             // Term3
- Sigma(k) * zp_a * zp_w                               // Term4

f_output = f_src0 * f_src1

(i_output-zero_point_output)*scale_output =((i_src0 - zero_point_0) * scale_0) * ((i_src1 - zero_point_1) * scale_1)
    
//经过变换
i_output = ((i_src0 - zero_point_0)*(i_src1 - zero_point_1)) * (scale_0*scale_1/scale_output) + zero_point_output

i_output = i_src0 *i_src1 * (scale_0*scale_1/scale_output)

std::pair<int32_t, int32_t> GetFixedPointMultiplierShift(double double_multiplier) {
    int32_t significand, exponent;
    if (double_multiplier == 0.) {
        significand = 0;
        exponent = 0;
        return std::make_pair(significand, exponent);
      }
    
    //获得尾数和阶数
    double significand_d = std::frexp(double_multiplier, &exponent);
    
    
    significand_d = std::round(significand_d * (1ll << 31));
    auto significand_int64 = static_cast<int64_t>(significand_d);
    ICHECK_LE(significand_int64, (1ll << 31));
    if (significand_int64 == (1ll << 31)) {
        significand_int64 /= 2;
        ++exponent;
    } //如果乘积的结果等于 (1ll << 31)，即超出了 32 位有符号整数的范围，则进行一次右移操作，并将指数部分加 1，以保证结果可以正确表示。
    ICHECK_LE(significand_int64, std::numeric_limits<int32_t>::max());
    significand = static_cast<int32_t>(significand_int64);
    return std::make_pair(significand, exponent);
    }

f_output = f_src0 * f_src1  //原始浮点计算
    
(i_output-zero_point_output)*scale_output =((i_src0 - zero_point_0) * scale_0) + ((i_src1 - zero_point_1) * scale_1)
    
//转换
i_output=((i_src0 - zero_point_0) * scale_0/scale_output) + ((i_src1 - zero_point_1) * scale_1/scale_output)+zero_point_output

import torch
import numpy as np

def get_tanh_lut( data,scale):


	data=data*scale # dequant   
	lut=data.tanh() #activation func
	lut=lut/scale # quant

	lut=lut.round()
	lut=torch.clamp(lut,-128,127)
	lut.to(torch.int)

	return lut

range_low=-128
range_high=127
scale = 0.2
data=torch.tensor(np.linspace(range_low,range_high,range_high-range_low+1))

lut = get_tanh_lut(data,scale)

int8 get_tanh_result(int8 input){
	# 这里偏移 128 是因为存放查找表时是-128，-127，...，126，127 这样的顺序， 对应的索引是 0,1,2，...，需要偏移 128
	return TANH_LUT[input+128]; 
}