Reconstruction-by-inpainting-for-visual-anomaly-detection icon indicating copy to clipboard operation
Reconstruction-by-inpainting-for-visual-anomaly-detection copied to clipboard

convert model to onnxruntime

Open vuhungtvt2018 opened this issue 2 years ago • 1 comments

how do you convert my model to onnxruntime? please help me to resolve this problem!

vuhungtvt2018 avatar Dec 13 '22 09:12 vuhungtvt2018

from models.unet import UNet import torch use_cuda = torch.cuda.is_available() device = torch.device('cuda' if use_cuda else 'cpu') model = UNet().to(device) model_path = r"D:\Work\doc\RIAD\best.pt" export_path = r"D:\Work\doc\RIAD\RIAD_model_320.onnx" check_point = torch.load(model_path, map_location=device) model.load_state_dict(check_point['model']) inputs = torch.randn(1, 3, 320, 320)

if name == 'main': torch.onnx.export( model, inputs, export_path, opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None )

bayueyaoshao avatar May 08 '24 02:05 bayueyaoshao