Add onnx export

[Luni-4]

This fixes #71

I used this command with InpaintingModel_best_gen.pth weight and config_ZITS_places2.yml file

python single_image_test.py --checkpoints ./checkpoints/ --config_file config_list/config_ZITS_places2.yml --onnx --save_path ./

and placed these weights StructureUpsampling.pth and best_transformer_places2.pth inside ckpt folder

There are some errors in ONNX export function though, can you help me to understand these errors?

BaseInpaintingTrainingModule init called
Loading InpaintingModel StructureUpsampling...
Loading trained transformer...
Loading InpaintingModel generator...
Warnning: There is no previous optimizer found. An initialized optimizer will be used.
Warnning: There is no previous optimizer found. An initialized optimizer will be used.
Traceback (most recent call last):
  File "/home/ZITS_inpainting/single_image_test.py", line 356, in <module>
    export_onnx(model, args.save_path)
  File "/home/ZITS_inpainting/single_image_test.py", line 297, in export_onnx
    torch.onnx.export(model,         # model being run
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 506, in export
    _export(
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 1525, in _export
    with exporter_context(model, training, verbose):
  File "/usr/lib/python3.10/contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 178, in exporter_context
    with select_model_mode_for_export(
  File "/usr/lib/python3.10/contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 139, in disable_apex_o2_state_d
ict_hook
    for module in model.modules():
AttributeError: 'ZITS' object has no attribute 'modules'

[DQiaole]

Hi, because 'ZITS' object is not a subclass of 'nn.Module'. All models are defined within 'ZITS.inpaint_model', where you can find our str_encoder, generator, structure_upsample, and transformer. You may need export these as ONNX format. Besides, our forward process of these submodels can be found in 'test' function of 'single_image_test.py'.

[bigmover]

This fixes #71

I used this command with InpaintingModel_best_gen.pth weight and config_ZITS_places2.yml file

python single_image_test.py --checkpoints ./checkpoints/ --config_file config_list/config_ZITS_places2.yml --onnx --save_path ./

and placed these weights StructureUpsampling.pth and best_transformer_places2.pth inside ckpt folder

There are some errors in ONNX export function though, can you help me to understand these errors?

BaseInpaintingTrainingModule init called
Loading InpaintingModel StructureUpsampling...
Loading trained transformer...
Loading InpaintingModel generator...
Warnning: There is no previous optimizer found. An initialized optimizer will be used.
Warnning: There is no previous optimizer found. An initialized optimizer will be used.
Traceback (most recent call last):
  File "/home/ZITS_inpainting/single_image_test.py", line 356, in <module>
    export_onnx(model, args.save_path)
  File "/home/ZITS_inpainting/single_image_test.py", line 297, in export_onnx
    torch.onnx.export(model,         # model being run
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 506, in export
    _export(
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 1525, in _export
    with exporter_context(model, training, verbose):
  File "/usr/lib/python3.10/contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 178, in exporter_context
    with select_model_mode_for_export(
  File "/usr/lib/python3.10/contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "/home/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 139, in disable_apex_o2_state_d
ict_hook
    for module in model.modules():
AttributeError: 'ZITS' object has no attribute 'modules'

perhaps can't. you succeed in converting the models？somemodule may fail
RuntimeError: Failed to export an ONNX attribute 'onnx::Sub', since it's not constant, please try to make things (e.g., kernel size) static if possible