Save the graph as SVG format. Args: output_file (str): a SVG file inputs: optional inputs given to the model. If given, the inputs will be used to run the graph to record shape of every tensor. The shape information will be
(self, output_file, inputs=None)
| 241 | f.write(self._init_net.SerializeToString()) |
| 242 | |
| 243 | def save_graph(self, output_file, inputs=None): |
| 244 | """ |
| 245 | Save the graph as SVG format. |
| 246 | |
| 247 | Args: |
| 248 | output_file (str): a SVG file |
| 249 | inputs: optional inputs given to the model. |
| 250 | If given, the inputs will be used to run the graph to record |
| 251 | shape of every tensor. The shape information will be |
| 252 | saved together with the graph. |
| 253 | """ |
| 254 | from .caffe2_export import run_and_save_graph |
| 255 | |
| 256 | if inputs is None: |
| 257 | save_graph(self._predict_net, output_file, op_only=False) |
| 258 | else: |
| 259 | size_divisibility = get_pb_arg_vali(self._predict_net, "size_divisibility", 0) |
| 260 | device = get_pb_arg_vals(self._predict_net, "device", b"cpu").decode("ascii") |
| 261 | inputs = convert_batched_inputs_to_c2_format(inputs, size_divisibility, device) |
| 262 | inputs = [x.cpu().numpy() for x in inputs] |
| 263 | run_and_save_graph(self._predict_net, self._init_net, inputs, output_file) |
| 264 | |
| 265 | @staticmethod |
| 266 | def load_protobuf(dir): |