Tensorrt自定义算子实现步骤
这篇文章主要是用于记录tensorrt自定义算子的过程.这里采用的是torch2trt直接转换的方式.第一步在模型中造一个自定义的算子:import torchfrom torch.nn.modules import Moduledef do_custom_Relu(x): return torch.clamp(x,min=0.0)def do_custom_op_shift(x,bias): x = x.cpu().detach().numpy() x = x + b
原创
2020-09-21 17:49:00 ·
5575 阅读 ·
1 评论