ONNX/TorchScript convertion about gpen

[jnulzl]

Hello!

Wanted to know if anyone managed to convert model with custom layers from stylegan2 to onnx/torchscript serialization format?

Convert to torchscript demo:

In gpen_model.py, you should

1、add export_torchscript attribute in __init__ of class Generator:

class Generator(nn.Module):
    def __init__(
        ......
        )
        ......      
        self.n_latent = self.log_size * 2 - 2
        self.export_torchscript = False ## add export_torchscript attribute 

    def make_noise(self):
         .....

2、add return image in forward of class Generator:

      if return_latents:
          return image, latent
      elif self.export_torchscript: ##add by jnulzl 
          return image              ##add by jnulzl 
      else:
          return image, None

3、In the end of gpen_model.py

if __name__ == '__main__':
    cpu_model = FullGenerator(512, 512, 8, 2).to(torch.device('cpu'))
    cpu_model.generator.export_torchscript = True
    data = torch.rand(1,3,512,512)
    _ = cpu_model(data)
    traced_cpu= torch.jit.trace(cpu_model, data)
    torch.jit.save(traced_cpu, "FullGenerator_cpu.pt")

[yangxy]

I will try it later.

[cyinen]

I will try it later.

Hello,have you convert the stylegan2 to onnx/torchscript finally?

[zimenglan-sysu-512]

hi @jnulzl what the verson of onnx u install is?