Add code for compiling model to torchscript
I want to be able to export it to other formats too in the future.
This commit is contained in:
parent
585b05e66b
commit
343af70a8d
17
codes/distill_torchscript.py
Normal file
17
codes/distill_torchscript.py
Normal file
|
@ -0,0 +1,17 @@
|
||||||
|
import argparse
|
||||||
|
import options.options as option
|
||||||
|
from models.networks import define_G
|
||||||
|
import torch
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('-opt', type=str, help='Path to options YMAL file.', default='options/test/test_ESRGAN_adrianna_full.yml')
|
||||||
|
opt = option.parse(parser.parse_args().opt, is_train=False)
|
||||||
|
opt = option.dict_to_nonedict(opt)
|
||||||
|
netG = define_G(opt)
|
||||||
|
|
||||||
|
print("Tracing generator network..")
|
||||||
|
dummyInput = torch.rand(1, 3, 8, 8)
|
||||||
|
traced_netG = torch.jit.trace(netG, dummyInput)
|
||||||
|
traced_netG.save('../results/traced_generator.zip')
|
||||||
|
print(traced_netG)
|
Loading…
Reference in New Issue
Block a user