pnnx 转ncnn，ncnn 推理的维度设置的in 和 torch.jit.trace inputshape 的维度不一样，有python 和 CPP 验证

[zuowanbushiwo]

detail | 详细描述 | 詳細な説明

Python 代码如下：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.enhance = nn.LSTM(input_size=512, hidden_size=128,  num_layers=1, dropout=0.0, batch_first=False)

    def forward(self, x, rnn_hc):
        hi =  rnn_hc[:1,:,:]
        ci =  rnn_hc[1:,:,:]
        out,(ho,co)  = self.enhance(x,(hi,ci))
        rnn_out_hc = torch.cat((ho, co), dim=0)
        return  out, rnn_out_hc   

def test1():
    rnn_hc = torch.zeros(2,1,128)
    input  = torch.randn([1,1,512])
    net = Model()
    out, rnn_hc = net(input,rnn_hc)
    print(out.shape)
    mod = torch.jit.trace(net,(input,rnn_hc)).eval()
    mod.save("test_ltsm.pt")
    import os
    os.system("~/pnnx/pnnx-20221116-ubuntu/pnnx  test_ltsm.pt  inputshape=[1,1,512],[2,1,128]")

if __name__ == '__main__':
    test1()

CPP 代码

CMakeLists.txt

cmake_minimum_required(VERSION 3.13)
project(test_ncnn)
set(CMAKE_CXX_STANDARD 14)
set(ncnn_DIR "/home/yang/ncnn_dcrn/ncnn-20221128-ubuntu-1804-shared/lib/cmake/ncnn" CACHE PATH "Directory that contains ncnnConfig.cmake")
find_package(ncnn REQUIRED)

set(SOURCE_FILES  main.cpp )
add_executable(test_ncnn ${SOURCE_FILES})
target_link_libraries(test_ncnn  ncnn)

#include <iostream>
#include <net.h>
int main()
{
    ncnn::Net  dcrn_net;
    dcrn_net.load_param("/home/yang/ncnn_cpp_test/test_ltsm.ncnn.param");
    dcrn_net.load_model("/home/yang/ncnn_cpp_test/test_ltsm.ncnn.bin");

    ncnn::Mat in0;
    in0.create(1, 1, 512);
    in0.fill(0.0f);

    ncnn::Mat in1;
    in1.create(128, 1, 2);
    //in1.create(2, 1, 128);    //如果打开这个地方就会崩溃，但是trace 是的inputshape 是 [2,1,128]
    in1.fill(0.0f);

    ncnn::Extractor ex1 = dcrn_net.create_extractor();
    ex1.set_num_threads(1);

    ex1.input("in0", in0);
    ex1.input("in1", in1);

    ncnn::Mat out2;
    ex1.extract("2", out2);
    printf("2: %d %d %d \n", out2.c,out2.h,out2.w);    

    ncnn::Mat out3;
    ex1.extract("3", out3);
    printf("3: %d %d %d \n", out3.c,out3.h,out3.w); 

    ncnn::Mat out0;
    ex1.extract("out0", out0);
    printf("out0: %d %d %d \n", out0.c,out0.h,out0.w);

    ncnn::Mat out1;
    ex1.extract("out1", out1);
    printf("out1: %d %d %d \n", out1.c,out1.h,out1.w);

}

如果是 in1.create(128, 1, 2); 这时输出都正常, 也是想要的维度:

2: 1 1 128
3: 1 1 128
out0: 1 1 128
out1: 2 1 128

如果是 in1.create(2, 1, 128); 推理时就会出现崩溃；

2: 1 1 2
3: 127 1 2
段错误

我觉得这个地方有点转不过弯来，为什么是 in1 是 (128, 1, 2)， 我开始一直使用(2, 1, 128) 就莫名其妙的崩溃，后来debug 到这个地方，再乱猜的，竟然对了。而且python 推理的时候也是(2,1,128).

def test_inference():
    torch.manual_seed(0)
    in0 = torch.rand(1, 1, 512, dtype=torch.float)
    in1 = torch.rand(2, 1, 128, dtype=torch.float)
    out = []

    with ncnn.Net() as net:
         net.load_param("test_ltsm.ncnn.param")
         net.load_model("test_ltsm.ncnn.bin")

         with net.create_extractor() as ex:
            ex.input("in0", ncnn.Mat(in0.squeeze(1).numpy()).clone())
            ex.input("in1", ncnn.Mat(in1.squeeze(1).numpy()).clone())

            _, out0 = ex.extract("out0")
            out.append(torch.from_numpy(np.array(out0)).unsqueeze(1))
            _, out1 = ex.extract("out1")
            out.append(torch.from_numpy(np.array(out1)).unsqueeze(1))

    if len(out) == 1:
        return out[0]
    else:
        return tuple(out)

test_ltsm.ncnn.param

7767517
5 8
Input                    in0                      0 1 in0
Input                    in1                      0 1 in1
Slice                    tensor_split_0           1 2 in1 2 3 -23300=2,1,-233 1=0
LSTM                     lstm_0                   3 3 in0 2 3 out0 5 6 0=128 1=262144 2=0 3=128
Concat                   cat_0                    2 1 5 6 out1 0=0

[deepage]

如果你去看ncnn::Mat的构造函数就能知道，你设定的in1.create(2, 1, 128)实际对应到pytorch的维度为[128, 1, 2]，因为ncnn::Mat的构造函数列表顺序是(w, h, d, c)这样的，有多个重载。

[zuowanbushiwo]

@deepage 感谢回复，那in0 是什么情况， 我是不是应该写成in0.create(512,1,1)? 没有报错是因为，c 和 h 是1的原因？

[deepage]

in0确实应该写成in0.create(512,1,1)，没有因为这个报错，可能是并不需要它维度上有越界吧。具体需要看这个输入参与的运算，我不会LSTM。而第二个输入会出错，是因为hidden_size设定成128了，这样维度就对不上了。

[zuowanbushiwo]

@deepage 非常感谢