Reconstruction-by-inpainting-for-visual-anomaly-detection
Reconstruction-by-inpainting-for-visual-anomaly-detection copied to clipboard
convert model to onnxruntime
how do you convert my model to onnxruntime? please help me to resolve this problem!
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 )