How to make symbolic function at this torch.augtograd.Function

I tried to export FastFCN model to onnx however it failed since…

RuntimeError: ONNX export failed: Couldn't export Python operator ScaledL2

Defined at:
/home/im/FastFCN/encoding/nn/encoding.py(37): forward
/home/im/anaconda3/envs/regseg/lib/python3.7/site-packages/torch/nn/modules/module.py(1039): _slow_forward
/home/im/anaconda3/envs/regseg/lib/python3.7/site-packages/torch/nn/modules/module.py(1051): _call_impl
/home/im/anaconda3/envs/regseg/lib/python3.7/site-packages/torch/nn/modules/container.py(139): forward
/home/im/anaconda3/envs/regseg/lib/python3.7/site-packages/torch/nn/modules/module.py(1039): _slow_forward
/home/im/anaconda3/envs/regseg/lib/python3.7/site-packages/torch/nn/modules/module.py(1051): _call_impl

The ScaledL2 function is like this.

class ScaledL2(autograd.Function):
    @staticmethod
    def forward(ctx, X, C, S):
        SL = (X.unsqueeze(2).expand(X.size(0), X.size(1), C.size(0), C.size(1)) -
              C.unsqueeze(0).unsqueeze(0)).pow_(2).sum(3).mul_(S.view(1, 1, C.size(0)))
        ctx.save_for_backward(X, C, S, SL)
        return SL

    @staticmethod
    def backward(ctx, GSL):
        X, C, S, SL = ctx.saved_variables

        tmp = (X.unsqueeze(2).expand(X.size(0), X.size(1), C.size(0), C.size(1)) - C.unsqueeze(0).unsqueeze(0)).mul_(
            (2 * GSL).mul_(S.view(1, 1, C.size(0))).unsqueeze(3)
        )

        GX = tmp.sum(2)
        GC = tmp.sum((0, 1)).mul_(-1)
        GS = SL.div(S.view(1, 1, C.size(0))).mul_(GSL).sum((0, 1))

        return GX, GC, GS

scaled_l2 = ScaledL2.apply
__all__ = [ 'scaled_l2']

I searched for solution and most said this will solve my problem. Therefore, I tried to make symbolic static method to ScaledL2, but I have no idea with torch._C
Please let me know how to make symbolic method to this ‘ScaledL2’

Sincerely.

1 Like