| 61 | Model("vgg16", torchvision.models.vgg.vgg16(), [1, 3, 224, 224])) |
| 62 | |
| 63 | def get_all_onnx_models(self, output_dir=default_gen_path): |
| 64 | if not os.path.exists(output_dir) or os.path.isfile(output_dir): |
| 65 | os.makedirs(output_dir) |
| 66 | for model in self.models: |
| 67 | output = "{}/{}.onnx".format(output_dir, model.name) |
| 68 | logging.debug( |
| 69 | "get model file from torchvision to: {}".format(output)) |
| 70 | net = model.torch_model |
| 71 | net.eval() |
| 72 | input_data = torch.randn(model.input_shape) |
| 73 | torch.onnx.export( |
| 74 | net, |
| 75 | input_data, |
| 76 | output, |
| 77 | export_params=True, |
| 78 | opset_version=12, |
| 79 | input_names=["data"], |
| 80 | output_names=["ret"], |
| 81 | ) |
| 82 | |
| 83 | def convert_to_mge(self, output_dir=default_gen_path): |
| 84 | for model in self.models: |