赞
踩
导出失败,因为pytorch缺少对于torch.fft.rfft傅里叶变换算子导出onnx的支持,因此无法导出
def oexport_onnx(img: np.ndarray, mask: np.ndarray, config_p: str, ckpt_p: str="./lama/configs/prediction/default.yaml", mod=8, device="cuda"): predict_config = OmegaConf.load(config_p) predict_config.model.path = ckpt_p # device = torch.device(predict_config.device) device = torch.device(device) train_config_path = os.path.join( predict_config.model.path, 'config.yaml') with open(train_config_path, 'r') as f: train_config = OmegaConf.create(yaml.safe_load(f)) train_config.training_model.predict_only = True train_config.visualizer.kind = 'noop' checkpoint_path = os.path.join( predict_config.model.path, 'models', predict_config.model.checkpoint ) model = load_checkpoint( train_config, checkpoint_path, strict=False, map_location='cpu') model.eval().cpu() dynamic_axes = { 'img' : {0 : 'batch_size', 2 : 'width', 3 : 'height'}, 'mask' : {0 : 'batch_size', 2 : 'width', 3 : 'height'}, 'output' : {0 : 'batch_size', 2 : 'width', 3 : 'height'}, } dummy_inputs = { "img": torch.randn((1, 3, 536, 800) ,dtype=torch.float), "mask": torch.randint(low=0, high=1, size=(1, 1, 536, 800), dtype=torch.float), } output_names = ['out_put'] torch.onnx.export(model, # model being run dummy_inputs, # model input (or a tuple for multiple inputs) "lama_model.onnx", # where to save the model (can be a file or file-like object) export_params=True, # store the trained parameter weights inside the model file opset_version=10, # the ONNX version to export the model to do_constant_folding=True, # whether to execute constant folding for optimization input_names = ['img', 'mask'], output_names=output_names, dynamic_axes=dynamic_axes )
aten::fft_rfftn' to ONNX opset version 10 is not supported
# lama\bin\saicinpainting\training\modules\ffc.py
# lama\saicinpainting\training\modules\ffc.py
无法解决,官方没有对于fft的支持。有一个解决办法是将该算子换成torch.atan(),但是效果据说不太好。
https://pytorch.org/docs/stable/onnx_supported_aten_ops.html
https://blog.csdn.net/xz1308579340/article/details/124908825
https://github.com/advimman/lama/issues/84
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。