StyleGANEX
StyleGANEX copied to clipboard
How to convert the model to onnx?
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.