一文玩转pytorch转onnx-tensorRT ——（C）测试onnx转tensorRT
说明
前文讲到了如何在onnx注册自定义层，以便onnx parsing时找到所对应的层。
在示例demo前，贴出遇到的问题
- 问题: ONNX and tensorRT: ERROR: Network must have at least one output
- 开始当作正常bug去处理的，找到了一些参考资料：1）既然没有输出，就标记一个输出; 2）input可能出现的问题
- 后来发现上面的方式对我的问题没有什么效果，最后经过巨量的微调后的debug，发现是因为在转onnx时强行修改了如下的代码，而upsample的修改导致结果不正确（ps：做研究，切莫随便捧别人臭脚，每一步都是有根据的才行，不然……）。实际上，这种写法不是对所有的pytorch（pytorch下的onnx）都是有效的，pytorch下的onnx有多个版本的。
import torch.onnx.symbolic
# Override Upsample's ONNX export from old opset if required (not needed for TRT 5.1+)
@torch.onnx.symbolic.parse_args('v', 'is')
def upsample_nearest2d(g, input, output_size):
height_scale = float(output_size[-2]) / input.type().sizes()[-2]
width_scale = float(output_size[-1]) / input.type().sizes()[-1]
return g.op("Upsample", input,
scales_f=(1, 1, height_scale, width_scale),
mode_s="nearest")
@torch.onnx.symbolic.parse_args('v', 'is', 'i')
def upsample_bilinear2d(g, input, output_size, align_corners):
height_scale = float(output_size[-2]) / input.type().sizes()[-2]
width_scale = float(output_size[-1]) / input.type().sizes()[-1]
return g.op("Upsample", input,
scales_f=(1, 1, height_scale, width_scale),
mode_s="linear")
torch.onnx.symbolic.upsample_bilinear2d = upsample_bilinear2d
torch.onnx.symbolic.upsample_nearest2d = upsample_nearest2d
补充
后来的其他实验遇到的一些小bug，记录如下：
- averagePool如果stride=1时（方便写block时会这么写），那么这一层没有实质意义的，但是转成的onnx，再转到tensorRT下时，会出现值检查的bug。这种情况可以直接在生成网络时，去掉这些层。当然，还有复杂一点点的，直接去掉可能导致一些逻辑的模型加载报错，所以可以用pad填充替代此层（当作占位符，用最简单的按序号加载参数时会用到）。
测试onnx转tensorRT
- python代码
class gnSOLE(nn.Module):
def __init__(self):
super().__init__()
out_channels = 2
self.lr = nn.LeakyReLU()
self.mp = nn.MaxPool2d((2, 2), stride=(2, 2))
self.gn = nn.GroupNorm(num_groups=2, num_channels=out_channels, eps=1e-5)
self.bn = nn.BatchNorm2d(out_channels)
for m in self.modules():
if isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
x1 = self.mp(x)
x1 = self.gn(self.lr(x1))
# x1 = upsample_to(x1, x)
p4 = x1 + x
return p4
def GN_ONNX():
model = gnSOLE()
model.to(torch.device('cuda'))
zero_input = torch.rand(1, 2, 8, 8).cuda()
# model(zero_input)
torch.onnx.export(model, zero_input, 'gn.onnx', verbose=True)
- c++代码。这儿包括编译方式，具体还可以参考“如何测试tensorRT添加自定义层”中的介绍。另外这里注释掉了common（tensorRT的demo中引用的头文件），注释代码仍旧保留，权当参考。
//nvcc -o gn test_onnx.cpp ../cuda/groupnorm.cu /usr/src/tensorrt/samples/common/logger.cpp
// -I/home/user/package/cub-1.8.0 -I/usr/src/tensorrt/samples/common/ -I./../cuda/ -L/usr/local/cuda/lib64
// -lcudart -lcuda -L/usr/local/lib/ -lnvonnxparser -L/usr/lib/x86_64-linux-gnu/ -lnvinfer
// -lnvparsers -lnvinfer_plugin
#include ] [--useDLACore=]\n" ;
std::cout << "--help Display help information\n";
std::cout
<< "--datadir Specify path to a data directory, overriding the default. This option can be used multiple times to add multiple directories. If no data directories are given, the default is to use (data/samples/mnist/, data/mnist/)"
<< std::endl;
std::cout
<< "--useDLACore=N Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, where n is the number of DLA engines on the platform."
<< std::endl;
std::cout << "--int8 Run in Int8 mode.\n";
std::cout << "--fp16 Run in FP16 mode." << std::endl;
}
int main(int argc, char **argv) {
auto sampleTest = gLogger.defineTest(gSampleName, argc, const_cast<const char **>(argv));
gLogger.reportTestStart(sampleTest);
// create a TensorRT model from the onnx model and serialize it to a stream
IHostMemory *trtModelStream{nullptr};
if (!onnxToTRTModel("/home/user/weight/gn.onnx", 1, trtModelStream))
std::cerr << "can not read onnx!";
assert(trtModelStream != nullptr);
uint8_t fileData[INPUT_H * INPUT_W];
// print an ascii representation
gLogInfo << "Input:\n"
// float data[INPUT_H * INPUT_W];
// for (int i = 0; i < INPUT_H * INPUT_W; i++)
// data[i] = 1.0 - float(fileData[i] / 255.0);
//
// // deserialize the engine
// IRuntime *runtime = createInferRuntime(gLogger);
// assert(runtime != nullptr);
// if (gArgs.useDLACore >= 0) {
// runtime->setDLACore(gArgs.useDLACore);
// }
//
// ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream->data(), trtModelStream->size(), nullptr);
// assert(engine != nullptr);
// trtModelStream->destroy();
// IExecutionContext *context = engine->createExecutionContext();
// assert(context != nullptr);
// // run inference
// float prob[OUTPUT_SIZE];
// doInference(*context, data, prob, 1);
//
// // destroy the engine
// context->destroy();
// engine->destroy();
// runtime->destroy();
//*********************************
// float val{0.0f};
// int idx{0};
//
// //Calculate Softmax
// float sum{0.0f};
// for (int i = 0; i < OUTPUT_SIZE; i++)
// {
// prob[i] = exp(prob[i]);
// sum += prob[i];
// }
//
// gLogInfo << "Output:\n";
// for (int i = 0; i < OUTPUT_SIZE; i++)
// {
// prob[i] /= sum;
// val = std::max(val, prob[i]);
// if (val == prob[i])
// idx = i;
//
// gLogInfo << " Prob " << i << " " << std::fixed << std::setw(5) << std::setprecision(4) << prob[i] << " "
// << "Class " << i << ": " << std::string(int(std::floor(prob[i] * 10 + 0.5f)), '*') << "\n";
// }
// gLogInfo << std::endl;
//
// bool pass{idx == num && val > 0.9f};
//
// return gLogger.reportTest(sampleTest, pass);
}