gmalivenko / pytorch2keras

PyTorch to Keras model convertor
https://pytorch2keras.readthedocs.io/en/latest/
MIT License
857 stars 143 forks source link

concatenate axis not getting changed when switching from NCHW mode to NHWC mode #131

Open yogeesh-alphaics opened 3 years ago

yogeesh-alphaics commented 3 years ago

Describe the bug I have a pth file for the SSD model which I wish to convert in Keras format, but I need to change the ordering (NCHW -> NHWC). In the end, there are 2 concat layers whose axis is 2 given the input is coming as x,y,z. Now if we change the ordering from nchw to nhwc , the input shape to concat also changes to [ x,z,y ]. But at this moment the axis of concat layer does not change from 2 to 1 and is resulting in error.

To Reproduce import torch import numpy as np from torchsummary import summary from torch.autograd import Variable

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") ssd_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd' , map_location=torch.device('cpu')) ssd_model.to("cpu") cl_model.eval() from pytorch2keras import pytorch_to_keras input_np = np.random.uniform(0,1,(1,3,300,300)) input_var = Variable(torch.FloatTensor(input_np)) kmodel = pytorch_to_keras(ssd_model , input_var , [(3,300,300)] , change_ordering = True , verbose = True) kmodel.summary() kmodel.save("./data/nhwc_model.h5" , save_format = "h5")

Expected behavior Ideally, if change_ordering is enabled, then concatenate axis should change accordingly. in this case axis should be 1 instead of 2

Logs ValueError: A Concatenate layer requires inputs with matching shapes except for the concat axis. Got inputs shapes: [(None, 5776, 4), (None, 2166, 4), (None, 600, 4), (None, 150, 4), (None, 36, 4), (None, 4, 4)]

Environment (please complete the following information):