pytorch1.8.0, onnx
F.relu6无法导出到onnx。应该是pytorch1.8.0没有定义relu6的导出规则,但是用nn.relu6是可以的。
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu = nn.ReLU6()
def forward(self, x):
# return torch.asinh(x)
return F.relu6(x)
model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'relu6.onnx', verbose=True, enable_onnx_checker=False, opset_version=11)
解决方案
- 将F.relu6 替换为F.hardtanh(x, 0, 6)
- 升级pytorch->1.10.0,实测是可以的。