Как преобразовать тензор Pytorch в тензор ONNX в пользовательском слое? - PullRequest
Новая
       108

Как преобразовать тензор Pytorch в тензор ONNX в пользовательском слое?

2 голосов
/ 14 апреля 2020

Я пытаюсь сделать 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'])
...