2

我正在尝试让 onnx 实现 pytorch 块。我用自定义转发功能制作了自定义块。

class MyConvBlockFunction(Function):

    @staticmethod
    def symbolic(g, input, conv1):
        from torch.onnx.symbolic_opset9 import _shape_as_tensor, _convolution, relu

        conv = _convolution(g, input, conv1.weight, False, 1, 1, 1, False, (), 1, None, None, None)
        output = relu(g, conv)

        return output

    @staticmethod
    def forward(self, input, conv1):
        conv = conv1(input)
        relu1 = nn.ReLU()
        res = relu1(conv)
        return res

class MyConvBlock(nn.Module):

    def __init__(self):
        super(MyConvBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 2, kernel_size=3, stride=1, padding=1, bias=False)
        self.relu = nn.ReLU()
        #self.weight = torch.tensor(self.conv1.weight, requires_grad=False)

    def forward(self, input):
        return MyConvBlockFunction.apply(input, self.conv1)

但是当我运行导出到 onnx 时，我得到一个错误，即权重张量是错误的张量。


weight_size = weight.type().sizes()
AttributeError: 'str' object has no attribute 'sizes'

我发现我的input张量是一个很好的onnx张量，而'conv.weight'是pytorch张量，而不是onnx张量。

input.1 defined in (%input.1 : Float(1, 3, 4, 4), %conv1.weight : Float(2, 3, 3, 3) = prim::Param()
2 defined in (%2 : Tensor = onnx::Constant[value=<Tensor>](), scope: MyConvBlock)

如何将权重发送到 onnx _convolution操作？

UPD：也许我的出口代码有问题？

torch.onnx.export(model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "convblock.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=OPSET_VER,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'])
4

0 回答 0