torch mobile with Adabins

https://pytorch.org/mobile/home/

torch mobile을 만드는 것을 시도하는 중이다. 코드는 adabins 공식 코드를 이용했다.

 

quantization은 옵션이고, torch.jit.script를 통해서 scriptmodule을 만들어야 optimize_for_mobile을 통해서 최적화를 할 수 있다. 

 

 

python = 3.8

torch = 1.9.0+cu111

 

에러1

torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/workspace/AdaBins_compression_test/models/unet_adaptive_bins.py", line 92
    def forward(self, x, **kwargs):
                          ~~~~~~~ <--- HERE
        unet_out = self.decoder(self.encoder(x), **kwargs)
        bin_widths_normed, range_attention_maps = self.adaptive_bins_layer(unet_out)

함수에 고정적이지 않은 argument 수가 적용될 수 없다고 한다. 꼭 있어야 하는 건 아니어서 제거했다.

 

에러2

RuntimeError: 
Module 'GenEfficientNet' has no attribute '_modules' :
  File "/workspace/AdaBins_compression_test/models/unet_adaptive_bins.py", line 68
    def forward(self, x):
        features = [x]
        for k, v in self.original_model._modules.items():
                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            if (k == 'blocks'):
                for ki, vi in v._modules.items():

참고: https://github.com/pytorch/pytorch/issues/28998

무슨 말인지 잘 모르겠지만 _modules가 iterator여서 알아서 ScriptModule을 생성해야하는 것으로 보인다. 누군가 사용 가능하게 업데이트 한 것 같은데 torch에 반영되지는 않은 듯 하다.

 

 

torch jit이 지원하지 못하는 부분이 많은 것 같다. 지금 efficientnet을 rwightman/gen-efficientnet-pytorch repo에서 가져와서 사용중인데 이 부분을 바꿔야 하나 싶다.