课程4：assignment check_equivalence

[nkzwchen]

如果把最后一层 softmax 去掉，原本可以正确运行的 check_equivalence 会报模型输出不相同的错误

IRModule 为

def mixed_pytorch_model(A, w4, b4, w6, b6):
    B = topi.nn.relu(A)
    B = topi.nn.pool2d(B, [2,2], [2,2], [1,1], padding=[0,0,0,0], pool_type="max")
    B = topi.nn.flatten(B)
    B = topi.matmul(B, w4, transp_b=True)

    B = topi.add(B, b4)
    B = topi.nn.relu(B)
    B = topi.matmul(B, w6, transp_b=True)
    
    B = topi.add(B, b6)
    # B = topi.nn.softmax(B)
    return B
def create_model_via_emit_te(batch_size, output_channel):
    bb = relax.BlockBuilder()
    x = relax.Var("x", input_shape, relax.DynTensorType(batch_size, "float32"))

    conv2d_weight = relax.const(weight_map["conv2d_weight"], "float32")
    conv2d_bias = relax.const(weight_map["conv2d_bias"].reshape(1, -1, 1, 1), "float32")
    linear0_weight = relax.const(weight_map["linear0_weight"], "float32")
    linear0_bias = relax.const(weight_map["linear0_bias"].reshape(1, -1), "float32")
    linear1_weight = relax.const(weight_map["linear1_weight"], "float32")
    linear1_bias = relax.const(weight_map["linear1_bias"].reshape(1, -1), "float32")
    
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit_te(tvm.topi.nn.conv2d, x, conv2d_weight, strides=1, padding=0, dilation=1)
            lv1 = bb.emit_te(tvm.topi.add, lv0, conv2d_bias)
            lv2 = bb.emit_te(
                mixed_pytorch_model, lv1,
                linear0_weight,
                linear0_bias,
                linear1_weight,
                linear1_bias)
            gv = bb.emit_output(lv2)
        bb.emit_func_output(gv)

    return bb.get()

pytorch model 为

def pytorch_model():
    list = []
    list.append(nn.Conv2d(in_channels=1, out_channels=output_channel, kernel_size=(3, 3), bias=True))
    list.append(nn.ReLU())
    list.append(nn.MaxPool2d(kernel_size=(2, 2)))
    list.append(nn.Flatten())
    list.append(nn.Linear(in_features=output_channel * 169, out_features=100, bias=True))
    list.append(nn.ReLU())
    list.append(nn.Linear(in_features=100, out_features=10, bias=True))
    # list.append(nn.Softmax(dim=1))

    model = nn.Sequential(*list).cpu()
    name_map = {
        "0.weight": "conv2d_weight",
        "0.bias": "conv2d_bias",
        "4.weight": "linear0_weight",
        "4.bias": "linear0_bias",
        "6.weight": "linear1_weight",
        "6.bias": "linear1_bias",
    }
    for name, param in model.named_parameters():
        param.data = torch.from_numpy(weight_map[name_map[name]]).cpu()
    return model

运行 check_equivalence 输出为

AssertionError: 
Not equal to tolerance rtol=0.0001, atol=1e-07

Mismatched elements: 1 / 40 (2.5%)
Max absolute difference: 1.1444092e-05
Max relative difference: 0.00012681
 x: array([[-6.399097e+00, -5.016352e+00, -5.534712e+00, -4.806492e+00,
        -4.344551e+00,  4.461435e+00, -1.965761e+00,  5.824947e+00,
         3.597243e+00,  1.564876e+01],...
 y: array([[-6.399096e+00, -5.016355e+00, -5.534713e+00, -4.806494e+00,
        -4.344555e+00,  4.461436e+00, -1.965763e+00,  5.824951e+00,
         3.597246e+00,  1.564877e+01],...

[Hzfengsy]

看起来只是精度问题，并且误差在可接受范围内