torch.bitwise_right_shift和bitwise_and导出到onnx算子的问题

代码中使用了bitpack来量化-反量化,其中用到了torch.bitwise_right_shift和bitwise_and函数,在导出ONNX模型时报错,特此记录。

1. Opset version选择

ONNX DOCUMENT可以看到,onnx直到opset 18才添加了对BitwiseAnd的支持。

2. 直接导出,报错

导出脚本:

torch.onnx.export(model, (x), "test.onnx", opset_version=18, verbose=True, input_names=['input'])

运行后:

    raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::bitwise_right_shift' to ONNX opset version 18 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

可以看到虽然ONNX 官方说了支持,但并没有与torch函数对应起来。

3. 添加函数支持

找到torch安装位置,打开torch/onnx/symbolic_opset18.py,添加下列代码:

@_onnx_symbolic("aten::bitwise_right_shift")
@_beartype.beartype
def bitwise_right_shift(g, input: torch.Value, other: torch.Value):
    return g.op("BitShift", input, other, direction_s='RIGHT')

@_onnx_symbolic("aten::bitwise_and")
@_beartype.beartype
def bitwise_and(g, input: torch.Value, other: torch.Value):
    return g.op("BitwiseAnd", input, other)

4. 运行

可能会遇到一些数据类型的报错,根据ONNX文档修改torch运行过程中涉及到tensor的dtype即可.

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
### 回答1: 好的,torch.bitwise_and()函数用于对两个张量进行按位与操作。以下是一个使用样例: ```python import torch a = torch.tensor([1, 2, 3], dtype=torch.uint8) b = torch.tensor([2, 3, 4], dtype=torch.uint8) result = torch.bitwise_and(a, b) print(result) ``` 输出结果为: ``` tensor([, 2, ], dtype=torch.uint8) ``` 这里将两个uint8类型的张量a和b进行按位与操作,得到了一个新的张量result。 ### 回答2: torch.bitwise_and()是PyTorch库中提供的一个函数,用来对两个输入的bit进行按位与操作。 以下是一个使用torch.bitwise_and()的示例: ``` import torch # 创建两个输入张量 a = torch.tensor([5, 3, 7], dtype=torch.uint8) b = torch.tensor([6, 3, 4], dtype=torch.uint8) # 对两个输入张量进行按位与操作 result = torch.bitwise_and(a, b) print(result) ``` 输出结果: ``` tensor([4, 3, 4], dtype=torch.uint8) ``` 在这个例子中,我们首先创建了两个输入张量a和b,它们的元素类型为torch.uint8,即8位无符号整数。 然后,我们使用torch.bitwise_and()函数对输入张量a和b进行按位与操作。在这个操作中,函数会对a和b的每个对应元素的bit进行与操作,得到一个新的张量result。 最后,我们打印输出结果result。由于按位与操作的定义是两个bit都为1时结果为1,否则为0,所以output张量的每个对应元素的bit都是通过将对应位置的a和b的bit进行与操作得到的。 注意:在进行bitwise_and操作时,请确保输入张量的元素类型是整数型,并且它们的形状相同。 ### 回答3: torch.bitwise_and()是PyTorch库中的一个函数,用于对两个张量进行按位与运算。下面是一个使用样例: ```python import torch # 创建两个输入张量 a = torch.tensor([5, 3, 1], dtype=torch.uint8) # [0101, 0011, 0001] b = torch.tensor([6, 3, 4], dtype=torch.uint8) # [0110, 0011, 0100] # 对两个张量进行按位与运算 c = torch.bitwise_and(a, b) # [0100, 0011, 0000] print(c) ``` 在这个样例中,我们首先导入了torch库。然后,我们创建了两个输入张量`a`和`b`,它们分别由三个8位元素组成。接下来,我们使用torch.bitwise_and()函数对两个张量进行按位与运算,并将结果保存在变量`c`中。最后,我们打印输出了`c`的结果。 在这个例子中,`a`和`b`的二进制表示分别为`[0101, 0011, 0001]`和`[0110, 0011, 0100]`,按位与运算的结果为`[0100, 0011, 0000]`。

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

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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值