[python] np.meshgrid(x, y) 와 torch.meshgrid(x, y) 는 다르다.

torchscript로 model을 저장하기 위해... numpy 연산을 없애야만 했다.

onnx 및 torch::jit::scriptModule로 변환을 하는 데 있어서 trace-based exporter가 numpy value를 잘 trace할 수 없다고 한다.

 

 

 

PyTorch models can be written using numpy manipulations, but this is not proper when we convert to the ONNX model. For the trace-based exporter, tracing treats the numpy values as the constant node, therefore it calculates the wrong result if we change the input. So the PyTorch model need implement using torch operators. For example, do not use numpy operators on numpy tensors:

출처:https://pytorch.org/docs/stable/onnx.html

torch.onnx — PyTorch 1.9.1 documentation

 

 

 

그래서 np.meshgrid(x, y)를 torch 연산으로 바꾸려고 찾아보니 torch에도 meshgrid 연산이 있어서 바로 가져다 썼다.

그랬더니 ㅇㅅㅇ...

np.meshgrid(x, y)의 결과 np array들의 shape과 torch.meshgrid(x, y)의 결과의 shape이 같지 않게 나오는 것이었다.

찾아보니 np.meshgrid에는 indexing option으로 'xy' 나 'ij'를 선택할 수 있는데, np.meshgrid의 default는 xy이고 torch.meshgrid는 np.meshgrid에 indexing='ij'를 준 방식으로 동작한다고 하였다.

np.meshgrid(x, y, indexing='ij') == torch.meshgrid(x, y)
np.meshgrid(x, y, indexing='xy') != torch.meshgrid(x, y)

정도로 해석할 수 있겠다.

 

아직까지 pip로 받을 수 있는 version의 pytorch에서는 torch.meshgrid의 indexing option을 지원하지 않고 있었는데

우연히도? 내가 이 문제를 겪기 18일전에 torch github에 xy 도 이제 지원을 한다는 commit message를 남겨둔 것을 발견했다.

https://github.com/pytorch/pytorch/pull/62724

implement "xy" indexing for torch.meshgrid by dagitses · Pull Request #62724 · pytorch/pytorch

 

여튼 source code build를 하면 간단히 해결될 문제였지만 귀찮아서 다른 방법이 있을거야 라고 생각했다.

노가다로 meshgrid 구현하려다가 잘 구현만 해놓으면 또 scriptModule로의 변환과정에서 지원안하는 타입이라고 떽떽대서 고생하던 중 번뜩 지나간 생각... 그 이름하여 transpose..

 

ㅇㅅㅇ..

그냥 

x_, y_ = torch.meshgrid(x, y)
x_ = torch.transpose(x)
y_ = torch.transpose(y)

를 해주면

x_, y_ = np.transpose(x, y)

와 같은 결과를 얻을 수 있다.

 

 

 

누군가는 나같은 삽질을 피하길...