StyleGANEX icon indicating copy to clipboard operation
StyleGANEX copied to clipboard

How to convert the model to onnx?

Open xiaocode opened this issue 1 month ago • 2 comments

from argparse import Namespace
from models.psp import pSp
import torch.nn as nn
import torch
import onnx

#Function to Convert to ONNX 
def Convert_ONNX(): 
    device = "cuda" if torch.cuda.is_available() else "cpu"
    ckpt_path = 'pretrained_models/styleganex_toonify_pixar.pt'
    ckpt = torch.load(ckpt_path, map_location='cpu')
    opts = ckpt['opts']
    opts['checkpoint_path'] = ckpt_path
    opts['device'] =  device
    opts = Namespace(**opts)
    torch_model = pSp(opts)
    torch_model.cpu()

    output_onnx = str("styleganex_toonify_pixar.onnx")

    # set the model to inference mode 
    torch_model.eval() 

    # The exported model will thus accept inputs of size [batch_size, 1, 224, 224] where batch_size can be variable.
    batch_size = 1 
    # Let's create a dummy input tensor
    channel = 3
    height = 224
    width = 224
    torch_input = torch.randn(batch_size, channel, height, width, requires_grad=True)

    dynamic_axes= {
        'input0': {0: 'batch', 2: 'height', 3: 'width'},
        'output0': {0: 'batch', 2: 'height', 3: 'width'}
    }

    # Export the model
    # """ 
    torch.onnx.export(
         torch_model,         # model being run 
         torch_input,       # model input (or a tuple for multiple inputs) 
         output_onnx,       # where to save the model  
         export_params=True,  # store the trained parameter weights inside the model file 
         opset_version=15,    # the ONNX version to export the model to 
         # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
         do_constant_folding=True,  # whether to execute constant folding for optimization
         input_names = ['input0'],   # the model's input names 
         output_names = ['output0'], # the model's output names 
         dynamic_axes = dynamic_axes)
    # """
    
    print(" ") 
    print('Model has been converted to ONNX')

    # Checks
    onnx_model = onnx.load(output_onnx)  # load onnx model
    onnx.checker.check_model(onnx_model)  # check onnx model

    print('ONNX export success, saved as %s' % output_onnx)


def main():
    Convert_ONNX()

if __name__ == "__main__":
    main()

When I run this code,it shows the error below:

torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of convolution for kernel of unknown shape. [Caused by the value '1865 defined in (%1865 : Float(*, *, *, *, strides=[401408, 784, 28, 1], requires_grad=1, device=cpu) = onnx::Reshape[allowzero=0](%1803, %1864), scope: models.psp.pSp::/models.stylegan2.model.Generator::decoder/models.stylegan2.model.StyledConv::conv1/models.stylegan2.model.ModulatedConv2d::conv # /home/yxy/github/StyleGANEX/models/stylegan2/model.py:297:0 )' (type 'Tensor') in the TorchScript graph. The containing node has kind 'onnx::Reshape'.]

I have searched some relative docs,It shows that we can not use dynamic shapes when convert to ONNX, but the doc in pytorch didn`t mention this.

xiaocode avatar May 07 '24 09:05 xiaocode