From 4ca6b1610f9fa65f8bd7d7c15059bfde18a2f02a Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 8 May 2018 00:56:35 -0700 Subject: Added data size and ONNX export options, FP16 inference is working --- test.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) (limited to 'test.py') diff --git a/test.py b/test.py index a9c8729..1effb08 100755 --- a/test.py +++ b/test.py @@ -26,6 +26,20 @@ webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.na for i, data in enumerate(dataset): if i >= opt.how_many: break + if opt.data_type == 16: + model.half() + data['label'] = data['label'].half() + data['inst'] = data['inst'].half() + elif opt.data_type == 8: + model.type(torch.uint8) + + if opt.export_onnx: + assert opt.export_onnx.endswith(".onnx"), "Export model file should end with .onnx" + if opt.verbose: + print(model) + generated = torch.onnx.export(model, [data['label'], data['inst']], + opt.export_onnx, verbose=True) + generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) -- cgit v1.2.3-70-g09d2