You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

tracing.py 31 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808
  1. # -*- coding: utf-8 -*-
  2. # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
  3. #
  4. # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
  5. #
  6. # Unless required by applicable law or agreed to in writing,
  7. # software distributed under the License is distributed on an
  8. # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  9. import collections
  10. import contextlib
  11. import functools
  12. import itertools
  13. import json
  14. import os
  15. import pickle
  16. import re
  17. import struct
  18. import sys
  19. from typing import Any
  20. import cv2
  21. import numpy as np
  22. from .. import tensor
  23. from ..core import _imperative_rt as rt
  24. from ..core._imperative_rt import GraphProfiler, GraphProfiler2, SerializationMetadata
  25. from ..core._imperative_rt.core2 import Tensor as RawTensor
  26. from ..core._imperative_rt.core2 import Trace, TraceError, name_tensor # skip_tracing,
  27. from ..core._imperative_rt.graph import _set_priority_to_id
  28. from ..core._imperative_rt.ops import (
  29. AssertEqual,
  30. CollectiveComm,
  31. ExternOpr,
  32. RemoteRecv,
  33. RemoteSend,
  34. set_jit_enabled,
  35. )
  36. from ..core._trace_option import set_symbolic_shape
  37. from ..core.tensor import megbrain_graph as G
  38. from ..logger import get_logger
  39. from ..utils import comp_graph_tools as cgtools
  40. from ..utils.naming import AutoNaming
  41. from ..utils.profiler import is_profiling
  42. from .dtr_config import DTRConfig
  43. from .graph_opt_config import GraphOptimizationConfig
  44. from .sublinear_memory_config import SublinearMemoryConfig
  45. logger = get_logger(__name__)
  46. def _input_node_use_static_shape():
  47. return os.environ.get("MEGENGINE_INPUT_NODE_USE_STATIC_SHAPE") is not None
  48. active_trace = None
  49. skip_tracing = False
  50. def is_tracing():
  51. if active_trace is None:
  52. return False
  53. else:
  54. return not skip_tracing
  55. @contextlib.contextmanager
  56. def exclude_from_trace():
  57. global skip_tracing
  58. if skip_tracing or (active_trace is None):
  59. yield
  60. return
  61. try:
  62. skip_tracing = True
  63. if active_trace is not None:
  64. active_trace._begin_excluded_region()
  65. yield
  66. if active_trace is not None:
  67. active_trace._end_excluded_region()
  68. finally:
  69. skip_tracing = False
  70. def array_comparator(lhs, rhs):
  71. return np.all(lhs == rhs)
  72. class trace:
  73. """Wraps a callable and provide:
  74. * tracing via :meth:`.trace` and :meth:`.dump`
  75. * accelerated evalutaion via :meth:`.__call__`
  76. Args:
  77. function: the function will be traced.
  78. symbolic: whether to apply symbolic execution for tracing. Default: False
  79. capture_as_const: capture global vars or closures as const value. Default: False
  80. record_only: if True, won't run even if call the function. Default: False
  81. sublinear_memory_config: configuration for sublinear memory optimization.
  82. If not None, it enables sublinear memory optimization with given setting.
  83. profiling: whether to profile compiled trace. Default: False
  84. opt_level: optimization level for compiling trace. Default: 2
  85. graph_opt_config: configuration for graph optimization. Default: None
  86. symbolic_shape: whether to use symbolic shape for tracing. Default: True
  87. """
  88. def __new__(cls, *args, **kwargs):
  89. if not args:
  90. return functools.partial(cls, **kwargs)
  91. return super().__new__(cls)
  92. def __init__(
  93. self,
  94. function,
  95. symbolic=False,
  96. capture_as_const=False,
  97. record_only=False,
  98. sublinear_memory_config: SublinearMemoryConfig = None,
  99. dtr_config: DTRConfig = None,
  100. profiling: bool = False,
  101. opt_level: int = 2,
  102. graph_opt_config: GraphOptimizationConfig = None,
  103. symbolic_shape: bool = True,
  104. ):
  105. self.__wrapped__ = function
  106. self._capture_as_const = capture_as_const or record_only
  107. self._arg_bindings = None
  108. self._kwarg_bindings = None
  109. self._output_bindings = None
  110. self._symbolic_shape = symbolic_shape
  111. self._graph_options = {
  112. "no_force_inplace": True,
  113. "graph_opt_level": opt_level,
  114. "seq_opt.enable_seq_comp_node_opt": False,
  115. }
  116. # prevent cyclic reference
  117. graph_options = self._graph_options
  118. if dtr_config is not None:
  119. graph_options["enable_dtr_memory_opt"] = True
  120. graph_options[
  121. "dtr_config.eviction_threshold"
  122. ] = dtr_config.eviction_threshold
  123. graph_options[
  124. "dtr_config.evictee_minimum_size"
  125. ] = dtr_config.evictee_minimum_size
  126. graph_options[
  127. "dtr_config.recomp_memory_factor"
  128. ] = dtr_config.recomp_memory_factor
  129. graph_options[
  130. "dtr_config.recomp_time_factor"
  131. ] = dtr_config.recomp_time_factor
  132. if graph_opt_config is not None:
  133. mapping = {None: 0, False: 1, True: 2}
  134. graph_options["graph_opt.jit_config.fuse_dimshuffle"] = mapping[
  135. graph_opt_config.jit_fuse_dimshuffle
  136. ]
  137. graph_options["graph_opt.jit_config.fuse_reduce"] = mapping[
  138. graph_opt_config.jit_fuse_reduce
  139. ]
  140. if sublinear_memory_config is not None:
  141. graph_options["enable_sublinear_memory_opt"] = True
  142. graph_options[
  143. "sublinear_mem_config.lb_memory_mb"
  144. ] = sublinear_memory_config.lb_memory_mb
  145. graph_options[
  146. "sublinear_mem_config.genetic_nr_iter"
  147. ] = sublinear_memory_config.genetic_nr_iter
  148. graph_options[
  149. "sublinear_mem_config.genetic_pool_size"
  150. ] = sublinear_memory_config.genetic_pool_size
  151. graph_options[
  152. "sublinear_mem_config.thresh_nr_try"
  153. ] = sublinear_memory_config.thresh_nr_try
  154. graph_options[
  155. "sublinear_mem_config.num_worker"
  156. ] = sublinear_memory_config.num_worker
  157. if int(os.getenv("MEGENGINE_INPLACE_UPDATE", "0")):
  158. graph_options["var_sanity_check_first_run"] = False
  159. def apply_options(options):
  160. for k, v in graph_options.items():
  161. words = k.split(".")
  162. suboptions = options
  163. for word in words[:-1]:
  164. suboptions = getattr(suboptions, word)
  165. setattr(suboptions, words[-1], v)
  166. self._trace = Trace()
  167. self._trace.symbolic = symbolic or record_only
  168. self._trace.capture_as_const = capture_as_const or record_only
  169. self._trace.no_exec = record_only
  170. self._trace.options_visitor = apply_options
  171. self._trace.profile = profiling
  172. self._trace.array_comparator = array_comparator
  173. self._trace.record_input_shapes = _input_node_use_static_shape()
  174. def __call__(self, *args, **kwargs):
  175. global active_trace
  176. symbolic_shape = None
  177. outputs = None
  178. try:
  179. active_trace = self
  180. self._trace.enter()
  181. if self._capture_as_const:
  182. self._process_inputs(*args, **kwargs)
  183. symbolic_shape = set_symbolic_shape(self._symbolic_shape)
  184. outputs = self.__wrapped__(*args, **kwargs)
  185. finally:
  186. handling_exc = sys.exc_info() != (None,) * 3
  187. active_trace = None
  188. if symbolic_shape is not None:
  189. symbolic_shape = set_symbolic_shape(symbolic_shape)
  190. assert symbolic_shape == self._symbolic_shape
  191. if self._capture_as_const and (outputs is not None):
  192. self._process_outputs(outputs)
  193. try:
  194. # may raise TraceError
  195. self._trace.exit()
  196. except TraceError:
  197. if not handling_exc:
  198. raise
  199. return outputs
  200. def _process_inputs(self, *args, **kwargs):
  201. for i, arg in enumerate(args):
  202. name_tensor("arg_{}".format(i), arg)
  203. # TODO: mark kwargs in order
  204. for k, kwarg in kwargs.items():
  205. if isinstance(kwarg, RawTensor):
  206. name_tensor("kwarg_{}".format(k), kwarg)
  207. if self._arg_bindings is None:
  208. self._arg_bindings = [
  209. ("arg_{}".format(i), arg._tuple_shape) for i, arg in enumerate(args)
  210. ]
  211. if self._kwarg_bindings is None:
  212. self._kwarg_bindings = {
  213. "kwarg_{}".format(k): (k, kwarg._tuple_shape)
  214. for k, kwarg in kwargs.items()
  215. if isinstance(kwarg, RawTensor)
  216. }
  217. def _process_outputs(self, outputs):
  218. if isinstance(outputs, RawTensor):
  219. outputs = [outputs]
  220. if isinstance(outputs, collections.abc.Mapping):
  221. output_names, outputs = zip(*sorted(outputs.items()))
  222. else:
  223. # output_names = ["output_{}".format(i) for i in range(len(outputs))]
  224. output_names = None
  225. self._output_names = output_names
  226. for i, output in enumerate(outputs):
  227. name_tensor("output_{}".format(i), output)
  228. if self._output_bindings is None:
  229. self._output_bindings = ["output_{}".format(i) for i in range(len(outputs))]
  230. def _begin_excluded_region(self):
  231. self._trace.begin_excluded_region()
  232. def _end_excluded_region(self):
  233. self._trace.end_excluded_region()
  234. def _make_feed(
  235. self,
  236. graph,
  237. outputs,
  238. input_data,
  239. repeat,
  240. silent,
  241. no_assert,
  242. maxerr,
  243. resize_input,
  244. input_transform,
  245. ):
  246. def auto_reformat_image(path, data, dst_shape):
  247. """reformat image to target shape
  248. :param data: image data as numpy array
  249. :param dst_shape: target shape
  250. """
  251. dim3_format = False # required input format does not contain batch
  252. hwc_format = False # required input format is NHWC
  253. if not dst_shape: # input tensor shape is not predefined
  254. if len(data.shape) == 2:
  255. chl = 1
  256. h = data.shape[0]
  257. w = data.shape[1]
  258. else:
  259. assert (
  260. len(data.shape) == 3
  261. ), "Input image must be of dimension 2 or 3"
  262. h, w, chl = data.shape
  263. dst_shape = (1, chl, h, w)
  264. if len(dst_shape) == 3:
  265. dst_shape = (1,) + dst_shape
  266. dim3_format = True
  267. assert len(dst_shape) == 4, "bad dst_shape: {}".format(dst_shape)
  268. chl = dst_shape[1]
  269. if chl in [1, 3]:
  270. n, c, h, w = dst_shape
  271. dst_shape = (n, h, w, c)
  272. else:
  273. chl = dst_shape[3]
  274. assert chl in [
  275. 1,
  276. 3,
  277. ], "can not infer input format from shape: {}".format(dst_shape)
  278. hwc_format = True
  279. # dst_shape has now been normalized to NHWC format
  280. if resize_input:
  281. h, w = dst_shape[1:3]
  282. data = cv2.resize(data, (w, h))
  283. logger.info("input {} resized to {}".format(path, data.shape))
  284. if chl == 1:
  285. data = cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)
  286. data = data[:, :, np.newaxis]
  287. assert data.ndim == 3
  288. data = data[np.newaxis]
  289. # data normalized to NHWC format
  290. if not hwc_format:
  291. data = np.transpose(data, (0, 3, 1, 2))
  292. if dim3_format:
  293. data = np.squeeze(data, 0)
  294. return data
  295. def read_input_data(dst_shape, dtype, path):
  296. def check_shape_equal(dst_shape, data_shape):
  297. if len(dst_shape):
  298. assert len(data_shape) == len(
  299. dst_shape
  300. ), "input/data shapes mismatch: {} vs {}".format(
  301. dst_shape, data_shape
  302. )
  303. if data_shape[1:] != dst_shape[1:]:
  304. logger.warning(
  305. "dst_shape is {}; data_shape is {}".format(
  306. dst_shape, data_shape
  307. )
  308. )
  309. if path.startswith("#"):
  310. assert not resize_input
  311. assert not input_transform
  312. spec = path
  313. m = re.match(
  314. r"^#rand\(([-0-9.]*)\s*,\s*([-0-9.]*)\s*(,[^\)]+)?\)$", spec
  315. )
  316. assert m, "bad spec {}".format(spec)
  317. rng_min = float(m.group(1))
  318. rng_max = float(m.group(2))
  319. if m.group(3):
  320. shape_str = m.group(3)
  321. try:
  322. shape = shape_str[1:].split(",")
  323. if shape[-1].strip() == "...":
  324. shape = shape[:-1]
  325. shape.extend(list(dst_shape[len(shape) :]))
  326. data_shape = tuple(map(int, shape))
  327. except ValueError as e:
  328. raise ValueError("bad spec {}: {}".format(spec, e.args))
  329. else:
  330. data_shape = dst_shape
  331. check_shape_equal(dst_shape, data_shape)
  332. return np.random.uniform(rng_min, rng_max, data_shape).astype(dtype)
  333. # try to load image
  334. data = cv2.imread(path, cv2.IMREAD_COLOR)
  335. if data is None:
  336. assert not resize_input
  337. data = np.load(path)
  338. assert isinstance(data, np.ndarray)
  339. else:
  340. # load image succeeds, so we expect input format is image format
  341. data = auto_reformat_image(path, data, dst_shape)
  342. data = np.repeat(data, repeat, axis=0)
  343. if repeat > 1:
  344. logger.info(
  345. "repeat input for {} times, data shape is {}".format(
  346. repeat, data.shape
  347. )
  348. )
  349. check_shape_equal(dst_shape, data.shape)
  350. if input_transform:
  351. data = eval(input_transform, {"data": data, "np": np})
  352. return data
  353. def gen_one_testcase(inputs, spec):
  354. paths = spec.split(";")
  355. if len(paths) != len(inputs):
  356. if len(paths) == 1 and paths[0].startswith("#"):
  357. paths = ["{}:{}".format(name, paths[0]) for name in inputs.keys()]
  358. assert len(paths) == len(
  359. inputs
  360. ), "required inputs: {}; data paths: {}".format(inputs.keys(), paths)
  361. if len(paths) == 1 and ":" not in paths[0]:
  362. paths[0] = next(iter(inputs.keys())) + ":" + paths[0]
  363. ret = {}
  364. for path in paths:
  365. var, path = path.split(":")
  366. ret[var] = read_input_data(inputs[var].shape, inputs[var].dtype, path)
  367. return ret
  368. inputs = cgtools.get_dep_vars(outputs, "Host2DeviceCopy")
  369. inputs = {i.name: i for i in inputs}
  370. if not no_assert:
  371. replace_varmap = {}
  372. inp_map = {}
  373. # replace var use InputNode
  374. for name, var in inputs.items():
  375. inp = G.InputNode(
  376. device="xpux", dtype=var.dtype, shape=var.shape, graph=graph
  377. )
  378. replace_varmap[var] = inp.outputs[0]._node
  379. inp_map[name] = inp
  380. new = cgtools.replace_vars(outputs, replace_varmap)
  381. if isinstance(new, rt.VarNode):
  382. new = list(new)
  383. output_nodes = [G.OutputNode(var) for var in new]
  384. func = graph.compile(*[node.outputs[0]._node for node in output_nodes])
  385. def make_dev_tensor(value, dtype=None, device=None):
  386. return tensor(value, dtype=dtype, device=device)._dev_tensor()
  387. def calculate(*args, **kwargs):
  388. output_val = []
  389. # set inputs value
  390. for name, var in inputs.items():
  391. val = kwargs.pop(name, None)
  392. assert val is not None, "miss input name{}".format(name)
  393. dev_tensor = make_dev_tensor(val, dtype=var.dtype, device="xpux")
  394. inp_map[name].set_value(dev_tensor)
  395. func.execute()
  396. for res in output_nodes:
  397. output_val.append(res.get_value().numpy())
  398. return output_val
  399. def expect_name(var):
  400. return "{}:expect".format(var.name)
  401. testcases = []
  402. np.set_printoptions(precision=2, threshold=4, suppress=True)
  403. data_list = []
  404. for item in input_data:
  405. if item.startswith("@"):
  406. with open(item[1:], "r") as f:
  407. data_list.extend(
  408. [line.rstrip() for line in f if line.rstrip() != ""]
  409. )
  410. else:
  411. data_list.append(item)
  412. for inp_spec in data_list:
  413. cur_testcase = gen_one_testcase(inputs, inp_spec)
  414. assert len(cur_testcase) == len(
  415. inputs
  416. ), "required inputs: {}; given data: {}".format(
  417. inputs.keys(), cur_testcase.keys()
  418. )
  419. if not no_assert:
  420. outputs_get = calculate(**cur_testcase)
  421. for var, val in zip(outputs, outputs_get):
  422. cur_testcase[expect_name(var)] = val
  423. logger.info(
  424. "generate test groundtruth: var={} shape={} range=({}, {})"
  425. " mean={} var={}".format(
  426. var,
  427. val.shape,
  428. val.min(),
  429. val.max(),
  430. np.mean(val),
  431. np.var(val),
  432. )
  433. )
  434. testcases.append(cur_testcase)
  435. logger.info(
  436. "add testcase: \n {}".format(
  437. "\n ".join(
  438. "{}: shape={} dtype={} range=({:.2f},{:.2f}) "
  439. "mean={:.2f} sd={:.2f}".format(
  440. k, v.shape, v.dtype, v.min(), v.max(), np.mean(v), np.std(v)
  441. )
  442. for k, v in sorted(cur_testcase.items())
  443. )
  444. )
  445. )
  446. if not no_assert:
  447. def expect_shp(var):
  448. ret = var.shape
  449. if ret:
  450. return ret
  451. return testcases[0][expect_name(var)].shape
  452. def assert_equal(expect, real, **kwargs):
  453. op = AssertEqual(**kwargs)
  454. (res,) = G.apply_normal_varnode(op, expect, real)
  455. return res._node
  456. verbose = not silent
  457. outputs_new = []
  458. for i in outputs:
  459. device = rt.CompNode("xpux")
  460. dtype = i.dtype
  461. name = expect_name(i)
  462. shape = expect_shp(i)
  463. # make expect output as one input of model.
  464. expect_get = rt.make_h2d(graph, device, dtype, shape, name)
  465. # insert assert opr to check expect and real.
  466. outputs_new.append(
  467. assert_equal(expect_get, i, verbose=verbose, maxerr=maxerr,)
  468. )
  469. inputs[expect_name(i)] = expect_get
  470. outputs = outputs_new
  471. return {"outputs": outputs, "testcases": testcases}
  472. def dump(
  473. self,
  474. file,
  475. *,
  476. arg_names=None,
  477. output_names=None,
  478. append=False,
  479. keep_var_name: int = 1,
  480. keep_opr_name: bool = False,
  481. keep_param_name: bool = False,
  482. keep_opr_priority: bool = False,
  483. strip_info_file=None,
  484. append_json=False,
  485. optimize_for_inference=True,
  486. user_info: Any = None,
  487. enable_metadata: bool = True,
  488. input_data=None,
  489. repeat=1,
  490. silent=False,
  491. no_assert=False,
  492. maxerr=1e-4,
  493. resize_input=False,
  494. input_transform=None,
  495. dump_format: str = None,
  496. **kwargs
  497. ):
  498. r"""Serializes trace to file system.
  499. Args:
  500. file: output file, could be file object or filename.
  501. arg_names: names of the input tensors in the traced function.
  502. output_names: names of the output tensors in the traced function,
  503. use the default name if not specified.
  504. append: whether output is appended to ``file``.
  505. Only works when ``file`` is str.
  506. keep_var_name: level for keeping variable names:
  507. * 0: none of the names are kept
  508. * 1: (default)keep names of output vars
  509. * 2: keep names of all (output and internal) vars
  510. keep_opr_name: whether to keep operator names.
  511. keep_param_name: whether to keep param names, so param values can be
  512. easily manipulated after loading model
  513. keep_opr_priority: whether to keep priority setting for operators
  514. strip_info_file: a string for path or a file handler. if is not None,
  515. then the dump information for code strip would be written to ``strip_info_file``
  516. append_json: will be check when `strip_info_file` is not None. if set
  517. true, the information for code strip will be append to strip_info_file.
  518. if set false, will rewrite strip_info_file
  519. optimize_for_inference: enbale optmizations,
  520. will skip all optimize options if this is False. Default: True
  521. user_info: any type object, which will be pickled to bytes.
  522. enable_metadata: whether to save metadata into output file.
  523. input_data: input test data and current network output would be used as groundtruth.
  524. The format is "var0:file0;var1:file1..." to specify data files for input vars.
  525. It can also be "#rand(min,max,shape...)" for generating random input data, for
  526. example, "#rand(0,255)", "#rand(0,255,1,3,224,224)" or "#rand(0, 255, 1, ...)"
  527. where `...` means the remaining part of the original shape. If the shape is not
  528. specified, the shape of corresponding input tensors in the network will be used.
  529. If there is only one input var, its name can be omitted. Each data file can either
  530. be an image which can be loaded by opencv, or a pickled numpy.ndarray. This option
  531. can be given multiple times to add multiple testcases. If you start the data
  532. with the letter @, the rest should be a filename, and each line in the file should
  533. be a single datum in the format described above. *NOTE* If `input_data` is not None,
  534. you can only use load-and-run to run the output file.
  535. repeat: how many times the input image is repeated. Useful when running benchmark for
  536. batch size other than one. Have no effect on randomly generated input data.
  537. silent: whether set verbose to False in assert_equal opr.
  538. no_assert: whether insert assert_equal opr to check result; this option is useful for
  539. benchmarking.
  540. maxerr: max error for assert_equal check during runtime.
  541. resize_input: whether resize input image to fit input var shape.
  542. input_transform: a python expression to transform the input data.
  543. Example: data / np.std(data)
  544. dump_format: using different dump formats. the open source MegEngine defaults to the FBS
  545. format. internal MegEngine have a choice of FBS and internal proprietary formats
  546. Keyword Arguments:
  547. * enable_io16xc32 --
  548. whether to use float16 for I/O between oprs and use
  549. float32 as internal computation precision. Note the output var would be
  550. changed to float16.
  551. * enable_ioc16 --
  552. whether to use float16 for both I/O and computation
  553. precision.
  554. * enable_hwcd4 --
  555. whether to use NHWCD4 data layout. This is faster on some
  556. OpenCL backend.
  557. * enable_nchw88 --
  558. whether to use NCHW88 data layout, currently
  559. used in X86 AVX backend.
  560. * enable_nchw44 --
  561. whether to use NCHW44 data layout, currently
  562. used in arm backend.
  563. * enable_nchw44_dot --
  564. whether to use NCHW44_dot data layout, currently
  565. used in armv8.2+dotprod backend.
  566. * enable_nchw4 --
  567. whether to use NCHW4 data layout, currently
  568. used in nvidia backend(based on cudnn).
  569. * enable_nchw32 --
  570. whether to use NCHW32 data layout, currently
  571. used in nvidia backend with tensorcore(based on cudnn).
  572. * enable_chwn4 --
  573. whether to use CHWN4 data layout, currently
  574. used in nvidia backend with tensorcore.
  575. * enable_nchw64 --
  576. whether to use NCHW64 data layout, used for fast int4
  577. support on Nvidia GPU.
  578. * enable_fuse_conv_bias_nonlinearity: whether to fuse conv+bias+nonlinearty
  579. into one opr.
  580. * enable_fuse_conv_bias_with_z: whether to fuse conv_bias with z
  581. input for inference on nvidia backend(this optimization pass will
  582. result in mismatch of the precision of output of training and
  583. inference)
  584. * enable_fuse_preprocess: whether to fuse astype\pad_channel\dimshuffle and
  585. etc opr
  586. """
  587. if not self._capture_as_const:
  588. raise ValueError(
  589. "you must specify capture_as_const=True at __init__ to use dump"
  590. )
  591. if self._output_names and output_names:
  592. raise TypeError(
  593. "cannot specify output_names when output is already in dict format"
  594. )
  595. if output_names and isinstance(output_names, str):
  596. output_names = (output_names,)
  597. if output_names and len(output_names) != len(self._output_bindings):
  598. raise ValueError(
  599. "wrong number of output_names, should be {} values".format(
  600. len(self._output_bindings)
  601. )
  602. )
  603. prefer_input_names = arg_names is not None
  604. if arg_names is None:
  605. arg_names = ["arg_%d" % i for i in range(len(self._arg_bindings))]
  606. if isinstance(arg_names, str):
  607. arg_names = (arg_names,)
  608. arg_names = [arg_name if arg_name is not None else "" for arg_name in arg_names]
  609. if arg_names and len(arg_names) != len(self._arg_bindings):
  610. raise ValueError(
  611. "wrong number of arg_names, should be {} values".format(
  612. len(self._arg_bindings)
  613. )
  614. )
  615. output_names = output_names or self._output_names
  616. if output_names is None:
  617. output_names = [""] * len(self._output_bindings)
  618. # output_names = ["output_{}".format(i) for i in range(len(self._output_bindings))]
  619. input_bindings = []
  620. def normalize_shape(shape):
  621. return (1,) if shape == () else shape
  622. for arg_name, (arg_id, arg_shape) in zip(arg_names, self._arg_bindings):
  623. input_bindings.append((arg_id, arg_name, normalize_shape(arg_shape)))
  624. for kwarg_id, (kwarg_name, kwarg_shape) in self._kwarg_bindings.items():
  625. input_bindings.append((kwarg_id, kwarg_name, normalize_shape(kwarg_shape)))
  626. graph = G.Graph()
  627. jit_enabled = set_jit_enabled(False)
  628. dest_vars = self._trace.dump(
  629. graph,
  630. input_bindings,
  631. [*zip(self._output_bindings, output_names)],
  632. prefer_input_names,
  633. )
  634. set_jit_enabled(jit_enabled)
  635. # dest_vars = [i._node for i in dest_vars]
  636. if input_data is not None:
  637. feeds = self._make_feed(
  638. graph,
  639. dest_vars,
  640. input_data,
  641. repeat,
  642. silent,
  643. no_assert,
  644. maxerr,
  645. resize_input,
  646. input_transform,
  647. )
  648. assert (
  649. isinstance(feeds, dict) and feeds["testcases"]
  650. ), "testcases can not be empty"
  651. dest_vars = feeds["outputs"]
  652. if optimize_for_inference:
  653. dest_vars, optimize_options = G.optimize_for_inference(dest_vars, **kwargs)
  654. dest_vars = [i._node for i in dest_vars]
  655. metadata = SerializationMetadata()
  656. if enable_metadata:
  657. metadata.user_info = pickle.dumps(user_info)
  658. metadata.is_valid = True
  659. metadata.graph_modified = False
  660. if optimize_for_inference:
  661. metadata.optimize_options = optimize_options
  662. if isinstance(file, str):
  663. permission = "wb" if append == False else "ab"
  664. file = open(file, permission)
  665. if keep_opr_priority:
  666. _set_priority_to_id(dest_vars)
  667. if input_data is not None:
  668. file.write(b"mgbtest0")
  669. file.write(struct.pack("I", len(feeds["testcases"])))
  670. dump_content, dump_info = G.dump_graph(
  671. dest_vars,
  672. keep_var_name=keep_var_name,
  673. keep_opr_name=keep_opr_name,
  674. keep_param_name=keep_param_name,
  675. keep_opr_priority=keep_opr_priority,
  676. strip_info_file=strip_info_file,
  677. append_json=append_json,
  678. metadata=metadata,
  679. dump_format=dump_format,
  680. )
  681. file.write(dump_content)
  682. if input_data is not None:
  683. inputs = cgtools.get_dep_vars(dest_vars, "Host2DeviceCopy")
  684. inputs = sorted((i.name, i.dtype) for i in inputs)
  685. def make_dev_tensor(value, dtype=None, device=None):
  686. return tensor(value, dtype=dtype, device=device)._dev_tensor()
  687. for testcase in feeds["testcases"]:
  688. assert isinstance(testcase, dict)
  689. cg = G.Graph()
  690. output_mgbvars = []
  691. for name, dtype in inputs:
  692. output_mgbvars.append(
  693. cg.make_const(
  694. make_dev_tensor(
  695. testcase.pop(name), dtype=dtype, device="cpux"
  696. )
  697. )
  698. )
  699. assert not testcase, "extra inputs provided in testcase: {}".format(
  700. testcase.keys()
  701. )
  702. dump_content, _ = G.dump_graph(
  703. output_mgbvars, strip_info_file=strip_info_file, append_json=True,
  704. )
  705. file.write(dump_content)
  706. return dump_info
  707. def get_profile(self):
  708. return json.loads(self._trace.get_profile())