(self, graph: BaseGraph, input_shapes: List[List[int]])
| 90 | json.dump(exports, file, indent=4) |
| 91 | |
| 92 | def prepare_model(self, graph: BaseGraph, input_shapes: List[List[int]]): |
| 93 | # trace model for exporting. |
| 94 | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| 95 | if all([var.value is not None for var in graph.inputs.values()]): |
| 96 | inputs = {var.name: convert_any_to_torch_tensor(var.value).to(device) for var in graph.inputs.values()} |
| 97 | elif all([var.shape is not None for var in graph.inputs.values()]): |
| 98 | inputs = {var.name: torch.randn(*var.shape, dtype=DataType.to_torch(var.dtype), device=device)\ |
| 99 | for var in graph.inputs.values()} |
| 100 | else: |
| 101 | assert len([input_shapes]) == len(graph.inputs), ( |
| 102 | 'must provide equal number of input shapes for caffe export without quantization') |
| 103 | # assume all fp32 type, because that's the usual case |
| 104 | inputs = [torch.randn(*shape, device=device) for shape in input_shapes] |
| 105 | tracer = TorchExecutor(graph, device=device) |
| 106 | tracer.tracing_operation_meta(inputs, list(graph.outputs.keys())) |
| 107 | |
| 108 | # build caffe protobuf |
| 109 | caffe_model = ppl_caffe_pb2.NetParameter() |
| 110 | caffe_model.name = graph._name if graph._name else 'PPQ Exported Caffe Model' |
| 111 | |
| 112 | # add caffe input info |
| 113 | for name, var in graph.inputs.items(): |
| 114 | caffe_model.input.append(name) |
| 115 | input_shape = ppl_caffe_pb2.BlobShape() |
| 116 | var.shape[0] = 1 |
| 117 | input_shape.dim.extend(var.shape) |
| 118 | caffe_model.input_shape.extend([input_shape]) |
| 119 | |
| 120 | # export op |
| 121 | for op in graph.topological_sort(): |
| 122 | if op.type not in caffe_export_map: |
| 123 | raise NotImplementedError( |
| 124 | f'{op.type} converted to Caffe OP is not supported in PPQ export parser yet') |
| 125 | |
| 126 | caffe_op = caffe_export_map[op.type](op) |
| 127 | assert isinstance(caffe_op, CaffeOpExporter) |
| 128 | |
| 129 | layer = caffe_op.parse() |
| 130 | if not isinstance(layer, (list, tuple)): |
| 131 | layer = [layer] |
| 132 | caffe_model.layer.extend(layer) |
| 133 | |
| 134 | caffe_model = optimize_for_export(caffe_model) |
| 135 | caffe_proto = deepcopy(caffe_model) |
| 136 | for layer in caffe_proto.layer: del layer.blobs[:] |
| 137 | |
| 138 | return caffe_model, caffe_proto |
| 139 | |
| 140 | def dump_to_file(self, caffe_model: ppl_caffe_pb2.NetParameter, |
| 141 | caffe_proto: ppl_caffe_pb2.NetParameter, file_path: str): |
no test coverage detected