converted DCGAN generator only accepts batches of 64 and doesn't work for batches of 1

[bilal2vec]

❓DCGAN generator only accepts batches of 64

System Information

• pytorch version: 1.3.1
• onnx-coreml version: 1.2
• coremltools version: 3.3

Code

model.py

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch

import numpy as np

class DCGAN_G(torch.nn.Module):
    def __init__(self):
        super(DCGAN_G, self).__init__()
        main = torch.nn.Sequential()

        # We need to know how many layers we will use at the beginning
        mult = 64 // 8

        # Start block
        # Z_size random numbers
        main.add_module('Start-ConvTranspose2d', torch.nn.ConvTranspose2d(100,
                                                                          64 * mult, kernel_size=4, stride=1, padding=0, bias=False))
        main.add_module('Start-BatchNorm2d',
                        torch.nn.BatchNorm2d(64 * mult))
        main.add_module('Start-ReLU', torch.nn.ReLU())
        # Size = (G_h_size * mult) x 4 x 4

        # Middle block (Done until we reach ? x image_size/2 x image_size/2)
        i = 1
        while mult > 1:
            main.add_module('Middle-ConvTranspose2d [%d]' % i, torch.nn.ConvTranspose2d(
                64 * mult, 64 * (mult//2), kernel_size=4, stride=2, padding=1, bias=False))
            main.add_module(
                'Middle-BatchNorm2d [%d]' % i, torch.nn.BatchNorm2d(64 * (mult//2)))
            main.add_module('Middle-ReLU [%d]' % i, torch.nn.ReLU())
            # Size = (G_h_size * (mult/(2*i))) x 8 x 8
            mult = mult // 2
            i += 1

        # End block
        # Size = G_h_size x image_size/2 x image_size/2
        main.add_module('End-ConvTranspose2d', torch.nn.ConvTranspose2d(64,
                                                                        3, kernel_size=4, stride=2, padding=1, bias=False))
        main.add_module('End-Tanh', torch.nn.Tanh())
        # Size = n_colors x image_size x image_size
        self.main = main

    def forward(self, input):

        # if isinstance(input.data, torch.cuda.FloatTensor) and 1 > 1:
        #     output = torch.nn.parallel.data_parallel(
        #         self.main, input, range(1))
        # else:
        output = self.main(input)
        return output

model = DCGAN_G()

# z_test = torch.FloatTensor(param.batch_size, param.z_size, 1, 1).normal_(0, 1)
z_test = torch.FloatTensor(1, 100, 1, 1).normal_(0, 1)

out = model(z_test)
print(out.shape)

torch.onnx.export(model,
                  z_test,
                  'gan.mlmodel',
                  input_names=["z"],
                  output_names=["img"])

broken_convert.py

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch

import numpy as np

import onnx_coreml

z_test = torch.FloatTensor(1, 100, 1, 1).normal_(0, 1)

mlmodel = onnx_coreml.convert(
    model='./moment-in-time.onnx', minimum_ios_deployment_target='13')

mlmodel.save('gan.mlmodel')

input_dict = {'z': z_test.numpy().astype(np.float32)}
pred_coreml = mlmodel.predict(input_dict, useCPUOnly=True)

print(pred_coreml['img'].shape)

working_convert.py

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch

import numpy as np

import onnx_coreml

z_test = torch.FloatTensor(64, 100, 1, 1).normal_(0, 1)

mlmodel = onnx_coreml.convert(
    model='./moment-in-time.onnx', minimum_ios_deployment_target='13')

mlmodel.save('gan.mlmodel')

input_dict = {'z': z_test.numpy().astype(np.float32)}
pred_coreml = mlmodel.predict(input_dict, useCPUOnly=True)

print(pred_coreml['img'].shape)

Reproduce

convert pytorch model to onnx

$ python model.py 

try converting onnx model to coreml and testing it with a batch size of 1, which doesn't work

$ python broken_convert.py

output:

WARNING:root:TensorFlow version 2.0.0 detected. Last version known to be fully compatible is 1.14.0 .
1/14: Converting Node Type ConvTranspose
2/14: Converting Node Type BatchNormalization
3/14: Converting Node Type Relu
4/14: Converting Node Type ConvTranspose
5/14: Converting Node Type BatchNormalization
6/14: Converting Node Type Relu
7/14: Converting Node Type ConvTranspose
8/14: Converting Node Type BatchNormalization
9/14: Converting Node Type Relu
10/14: Converting Node Type ConvTranspose
11/14: Converting Node Type BatchNormalization
12/14: Converting Node Type Relu
13/14: Converting Node Type ConvTranspose
14/14: Converting Node Type Tanh
Translation to CoreML spec completed. Now compiling the CoreML model.
Model Compilation done.
Traceback (most recent call last):
  File "convert.py", line 19, in <module>
    pred_coreml = mlmodel.predict(input_dict, useCPUOnly=True)
  File "/Users/bilal/miniconda3/envs/pytorch/lib/python3.7/site-packages/coremltools/models/model.py", line 334, in predict
    return self.__proxy__.predict(data, useCPUOnly)
RuntimeError: {
    NSLocalizedDescription = "Shape (1 x 100 x 1 x 1) was not in enumerated set of allowed shapes";
}

Try converting onnx model to coreml and testing it with a batch size of 64, which does work

$ python working_convert.py

output:

WARNING:root:TensorFlow version 2.0.0 detected. Last version known to be fully compatible is 1.14.0 .
1/14: Converting Node Type ConvTranspose
2/14: Converting Node Type BatchNormalization
3/14: Converting Node Type Relu
4/14: Converting Node Type ConvTranspose
5/14: Converting Node Type BatchNormalization
6/14: Converting Node Type Relu
7/14: Converting Node Type ConvTranspose
8/14: Converting Node Type BatchNormalization
9/14: Converting Node Type Relu
10/14: Converting Node Type ConvTranspose
11/14: Converting Node Type BatchNormalization
12/14: Converting Node Type Relu
13/14: Converting Node Type ConvTranspose
14/14: Converting Node Type Tanh
Translation to CoreML spec completed. Now compiling the CoreML model.
Model Compilation done.
(64, 3, 64, 64)

[bilal2vec]

Fixed, I was saving my onnx model as a .mlmodel file