【pytorch】——自定义一个算子并导出到onnx

pytorch, onnx

摘要:为了将自定义算子的参数,或者自己想要保存的参数序列化到onnx中。

code

import torch
import torch.nn as nn
from torch.autograd import Function
import onnx
import torch.onnx

class Requant_(Function):
    @staticmethod
    def forward(ctx, input, requant_scale, shift):               # ctx 必须要
        input = input.double() * requant_scale / 2**shift        # 为了等价于c中的移位操作。会存在int32溢出
        input = torch.floor(input).float()

        return torch.floor(input)
    
    @staticmethod
    def symbolic(g, *inputs):
        return g.op("Requant", inputs[0], scale_f=23.0, shift_i=8)

requant_ = Requant_.apply

class TinyNet(nn.Module):
    def __init__(self):
        super(TinyNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 1, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = x.view(-1)
        x = requant_(x, 5, 5)
        return x

net = TinyNet().cuda()
ipt = torch.ones(2,3,12,12).cuda()
torch.onnx.export(net, (ipt,), 'tinynet.onnx', opset_version=11, enable_onnx_checker=False)
print(onnx.load('tinynet.onnx'))

关键点:

  • 继承自torch.autograd
  • scale_f=23.0, shift_i=8,_f表示浮点数,_i表示整形int32类型

onnx 模型
在这里插入图片描述

总结

这种是在pytorch中新写一个op,并序列化到onnx中,另外一个想法是:如果修改已有op的onnx序列化,比如conv2d,upsample等。得到onnx模型中,还需要对onnx模型解析,在把onnx模型转换成自己想要的表达。

  • 0
    点赞
  • 16
    收藏
    觉得还不错? 一键收藏
  • 10
    评论
发布自定义的C++算子,需要完成以下步骤: 1. 编写C++代码:编写自定义算子的C++代码,包括前向计算和反向传播函数等。 2. 编写Python绑定代码:编写Python绑定代码,将C++算子封装为PyTorch模块,方便在Python中使用。 3. 编译C++代码:使用CMake或者Makefile等工具,将C++代码编译成共享库(.so或.dll文件)。 4. 使用PyTorch C++扩展API:使用PyTorch C++扩展API,将共享库加载到PyTorch中,并注册自定义算子。 5. 测试自定义算子:在Python中测试自定义算子是否能够正常工作。 这里提供一个简单的示例: 1. 编写C++代码 ```cpp #include <torch/extension.h> torch::Tensor my_add_forward(const torch::Tensor& input1, const torch::Tensor& input2) { return input1 + input2; } std::vector<torch::Tensor> my_add_backward(const torch::Tensor& grad_output) { return {grad_output, grad_output}; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &my_add_forward, "MyAdd forward"); m.def("backward", &my_add_backward, "MyAdd backward"); } ``` 2. 编写Python绑定代码 ```python import torch my_add = torch.utils.cpp_extension.load(name='my_add', sources=['my_add.cpp'], verbose=True) def my_add_op(input1, input2): return my_add.forward(input1, input2) class MyAddFunction(torch.autograd.Function): @staticmethod def forward(ctx, input1, input2): output = my_add_op(input1, input2) ctx.save_for_backward(input1, input2) return output @staticmethod def backward(ctx, grad_output): input1, input2 = ctx.saved_tensors grad_input = my_add.backward(grad_output) return grad_input[0], grad_input[1] my_add_function = MyAddFunction.apply ``` 3. 编译C++代码 使用以下命令编译C++代码: ```sh g++ -o my_add.so -shared -fPIC my_add.cpp $(python3 -m pybind11 --includes) -I/path/to/torch/include -I/path/to/torch/include/torch/csrc/api/include -L/path/to/torch/lib -ltorch -lc10 ``` 4. 使用PyTorch C++扩展API ```cpp #include <torch/script.h> #include <iostream> int main() { torch::jit::script::Module module = torch::jit::load("model.pt"); module.to(torch::kCPU); std::string code = R"( def forward(x, y): return my_add_function(x, y) )"; torch::jit::script::Module new_module = module.clone(); new_module.define(code); // Test the new module torch::Tensor x = torch::ones({2, 3}); torch::Tensor y = torch::ones({2, 3}); torch::Tensor output = new_module.forward({x, y}).toTensor(); std::cout << output << std::endl; return 0; } ``` 5. 测试自定义算子 在Python中测试自定义算子: ```python import torch my_add = torch.utils.cpp_extension.load(name='my_add', sources=['my_add.cpp'], verbose=True) def my_add_op(input1, input2): return my_add.forward(input1, input2) class MyAddFunction(torch.autograd.Function): @staticmethod def forward(ctx, input1, input2): output = my_add_op(input1, input2) ctx.save_for_backward(input1, input2) return output @staticmethod def backward(ctx, grad_output): input1, input2 = ctx.saved_tensors grad_input = my_add.backward(grad_output) return grad_input[0], grad_input[1] my_add_function = MyAddFunction.apply # Test the custom operator x = torch.ones(2, 3, requires_grad=True) y = torch.ones(2, 3, requires_grad=True) z = my_add_function(x, y) z.sum().backward() print(x.grad) # tensor([[1., 1., 1.],[1., 1., 1.]]) print(y.grad) # tensor([[1., 1., 1.],[1., 1., 1.]]) ```

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 10
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值