| 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): |