| 26 | |
| 27 | |
| 28 | class AllModel: |
| 29 | models = [] |
| 30 | |
| 31 | # model src from onnx |
| 32 | def __init__(self): |
| 33 | # pytorch model |
| 34 | self.models.append( |
| 35 | Model( |
| 36 | "mobilenetv2", |
| 37 | torchvision.models.mobilenetv2.mobilenet_v2(), |
| 38 | [1, 3, 224, 224], |
| 39 | )) |
| 40 | self.models.append( |
| 41 | Model( |
| 42 | "efficientnetb0", |
| 43 | torchvision.models.efficientnet.efficientnet_b0(), |
| 44 | [1, 3, 256, 256], |
| 45 | )) |
| 46 | self.models.append( |
| 47 | Model( |
| 48 | "shufflenetv2", |
| 49 | torchvision.models.shufflenetv2.shufflenet_v2_x0_5(), |
| 50 | [1, 3, 224, 224], |
| 51 | )) |
| 52 | self.models.append( |
| 53 | Model("resnet18", torchvision.models.resnet.resnet18(), |
| 54 | [1, 3, 224, 224])) |
| 55 | self.models.append( |
| 56 | Model("resnet50", torchvision.models.resnet.resnet50(), |
| 57 | [1, 3, 224, 224])) |
| 58 | self.models.append( |
| 59 | Model("vgg11", torchvision.models.vgg.vgg11(), [1, 3, 224, 224])) |
| 60 | self.models.append( |
| 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: |
| 85 | input = "{}/{}.onnx".format(output_dir, model.name) |