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.

network.py 29 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823
  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 fnmatch
  11. import itertools
  12. import pickle
  13. import re
  14. from collections import OrderedDict
  15. from typing import Any, Dict, List, Optional, Sequence
  16. from ..core import _imperative_rt
  17. from ..core._imperative_rt import ComputingGraph, SerializationMetadata
  18. from ..core._trace_option import set_symbolic_shape as _set_symbolic_shape
  19. from ..core.tensor import megbrain_graph as G
  20. from ..logger import get_logger
  21. from .comp_graph_tools import get_dep_vars, get_opr_type, get_oprs_seq
  22. from .network_node import (
  23. ConstOpBase,
  24. Host2DeviceCopy,
  25. ImmutableTensor,
  26. NetworkNode,
  27. OpNode,
  28. VarNode,
  29. str_to_mge_class,
  30. )
  31. logger = get_logger(__name__)
  32. class Network:
  33. def __init__(self):
  34. self.input_vars = [] # input var of graph
  35. self._orig_inputs = []
  36. self.output_vars = [] # output var of graph
  37. self._orig_outputs = []
  38. self.all_oprs_map = OrderedDict() # _imperative_rt.graph.VarNode.id: VarNode
  39. self.all_vars_map = (
  40. OrderedDict()
  41. ) # _imperative_rt.graph.OperatorNode.id: OpNode
  42. self.graph = ComputingGraph()
  43. self._metadata = None
  44. @property
  45. def metadata(self):
  46. r"""Load metadata as a dict."""
  47. if not self._metadata.is_valid:
  48. logger.info("metadata is not valid!")
  49. return None
  50. ret = dict()
  51. try:
  52. user_info = pickle.loads(self._metadata.user_info)
  53. except: # pylint: disable=bare-except
  54. logger.warning(
  55. "can't parse user info by pickle, so return the original bytes object!"
  56. )
  57. user_info = self._metadata.user_info
  58. ret["user_info"] = user_info
  59. ret["graph_modified"] = self._metadata.graph_modified
  60. ret["optimized_for_inference"] = self._metadata.optimized_for_inference
  61. if ret["optimized_for_inference"]:
  62. ret.update(G.deserialize_infer_option(self._metadata.optimize_options))
  63. return ret
  64. @classmethod
  65. def load(cls, model_path: str, outspec: List[str] = None):
  66. r"""Loads a computing graph as a Network object.
  67. Args:
  68. model_path: file path of mge model.
  69. outspec: only load the subgraph with outspec as its endpoints.
  70. """
  71. self = cls()
  72. ret = G.load_graph(model_path)
  73. outputs, self._metadata = ret.output_vars_list, ret.metadata
  74. if outspec is not None:
  75. output_spec = outspec.copy()
  76. all_vars = get_dep_vars(outputs) + outputs
  77. new_outputs = {}
  78. for i in all_vars:
  79. if i.name in output_spec:
  80. new_outputs[i.name] = i
  81. output_spec.remove(i.name)
  82. assert len(output_spec) == 0, "Can not find {} in this model".format(
  83. output_spec
  84. )
  85. outputs = [new_outputs[i] for i in outspec]
  86. self._orig_outputs = outputs
  87. for x in self._orig_outputs:
  88. self.output_vars.append(self._get_var(x))
  89. self.add_dep_oprs()
  90. for x in self._orig_inputs:
  91. self.input_vars.append(self._get_var(x))
  92. self.graph = self._orig_outputs[0].graph
  93. return self
  94. def _compile(self):
  95. self.all_oprs_map = {}
  96. self.all_vars_map = {}
  97. for opr in self.all_oprs:
  98. if isinstance(opr, (ConstOpBase, Host2DeviceCopy)):
  99. opr.compile(self.graph)
  100. else:
  101. opr.compile()
  102. if opr.name is not None:
  103. opr._opr.name = opr.name
  104. self.all_oprs_map[opr._opr.id] = opr
  105. for o in opr.outputs:
  106. self.all_vars_map[o.var.id] = o
  107. def optimize_for_inference(self, dest_vars, **kwargs):
  108. r"""Applies optimize_for_inference pass for operator graph.
  109. Args:
  110. dest_vars: list of output vars in the operator graph
  111. Keyword Arguments:
  112. * enable_io16xc32 --
  113. whether to use float16 for I/O between oprs and use
  114. float32 as internal computation precision. Note the output var would be
  115. changed to float16.
  116. * enable_ioc16 --
  117. whether to use float16 for both I/O and computation
  118. precision.
  119. * enable_hwcd4 --
  120. whether to use NHWCD4 data layout. This is faster on some
  121. OpenCL backend.
  122. * enable_nchw88 --
  123. whether to use NCHW88 data layout, currently
  124. used in X86 AVX backend.
  125. * enable_nchw44 --
  126. whether to use NCHW44 data layout, currently
  127. used in arm backend.
  128. * enable_nchw44_dot --
  129. whether to use NCHW44_dot data layout, currently
  130. used in armv8.2+dotprod backend.
  131. * enable_nchw4 --
  132. whether to use NCHW4 data layout, currently
  133. used in nvidia backend(based on cudnn).
  134. * enable_nchw32 --
  135. whether to use NCHW32 data layout, currently
  136. used in nvidia backend with tensorcore(based on cudnn).
  137. * enable_chwn4 --
  138. whether to use CHWN4 data layout, currently
  139. used in nvidia backend with tensorcore.
  140. * enable_nchw64 --
  141. whether to use NCHW64 data layout, used for fast int4
  142. support on Nvidia GPU.
  143. * enable_fuse_conv_bias_nonlinearity: whether to fuse conv+bias+nonlinearty
  144. into one opr.
  145. * enable_fuse_conv_bias_with_z: whether to fuse conv_bias with z
  146. input for inference on nvidia backend(this optimization pass will
  147. result in mismatch of the precision of output of training and
  148. inference)
  149. """
  150. if not isinstance(dest_vars, Sequence):
  151. dest_vars = [dest_vars]
  152. dest_vars = list(G.VarNode(var.var) for var in dest_vars)
  153. new_vars = G.optimize_for_inference(dest_vars, **kwargs)
  154. return list(self._get_var(var) for var in new_vars)
  155. def dump(
  156. self,
  157. file,
  158. *,
  159. keep_var_name: int = 1,
  160. keep_opr_name: bool = False,
  161. keep_param_name: bool = False,
  162. keep_opr_priority: bool = False,
  163. strip_info_file=None,
  164. append_json=False,
  165. optimize_for_inference=True,
  166. append=False,
  167. user_info: Any = None,
  168. enable_metadata=True,
  169. **kwargs
  170. ):
  171. r"""Serializes graph to file.
  172. Args:
  173. file: output file, could be file object or filename.
  174. append: whether output is appended to ``file``.
  175. Only works when ``file`` is str.
  176. keep_var_name: level for keeping variable names:
  177. * 0: none of the names are kept
  178. * 1: (default)keep names of output vars
  179. * 2: keep names of all (output and internal) vars
  180. keep_opr_name: whether to keep operator names.
  181. keep_param_name: whether to keep param names, so param values can be
  182. easily manipulated after loading model
  183. keep_opr_priority: whether to keep priority setting for operators
  184. strip_info_file: a string for path or a file handler. if is not None,
  185. then the dump information for code strip would be written to ``strip_info_file``
  186. append_json: will be check when `strip_info_file` is not None. if set
  187. true, the information for code strip will be append to strip_info_file.
  188. if set false, will rewrite strip_info_file
  189. optimize_for_inference: enbale optmizations,
  190. will skip all optimize options if this is False. Default: True
  191. user_info: any type object, which will be pickled to bytes.
  192. enable_metadata: whether to save metadata into output file.
  193. See more detials in :meth:`~.trace.dump`.
  194. """
  195. def _set_var_name(var):
  196. graph_var = G.VarNode(var.var)
  197. graph_var.name = var.name
  198. return graph_var
  199. self._compile()
  200. out = list(map(_set_var_name, self.output_vars))
  201. if kwargs.pop("arg_names", False):
  202. logger.warning(
  203. '"arg_names" is not supported in Network.dump, rename input vars directly'
  204. )
  205. if kwargs.pop("output_names", False):
  206. logger.warning(
  207. '"output_names" is not supported in Network.dump, rename output vars directly'
  208. )
  209. if optimize_for_inference:
  210. out, optimize_options = G.optimize_for_inference(out, **kwargs)
  211. metadata = SerializationMetadata()
  212. if enable_metadata:
  213. metadata.is_valid = True
  214. metadata.graph_modified = True
  215. metadata.user_info = pickle.dumps(user_info)
  216. if optimize_for_inference:
  217. metadata.optimize_options = optimize_options
  218. G.set_priority_to_id([o._node if isinstance(o, G.VarNode) else o for o in out])
  219. dump_content, _ = G.dump_graph(
  220. out,
  221. keep_var_name=keep_var_name,
  222. keep_opr_name=keep_opr_name,
  223. keep_param_name=keep_param_name,
  224. keep_opr_priority=keep_opr_priority,
  225. strip_info_file=strip_info_file,
  226. append_json=append_json,
  227. metadata=metadata,
  228. )
  229. if isinstance(file, str):
  230. permission = "wb" if append == False else "ab"
  231. file = open(file, permission)
  232. file.write(dump_content)
  233. def make_const(self, data, name=None, device=None):
  234. r"""Makes an ImmutableTensor OpNode to provide a parameter for the network."""
  235. node = ImmutableTensor(data, name, device, self.graph)
  236. node.compile(self.graph)
  237. return node.outputs[0]
  238. def make_input_node(self, shape, dtype, name=None, device=None):
  239. r"""Makes a Host2DeviceCopy OpNode to provide an input varnode for the network."""
  240. node = Host2DeviceCopy(shape, dtype, name, device)
  241. node.compile(self.graph)
  242. return node.outputs[0]
  243. def add_output(self, *vars: VarNode):
  244. r"""Adds vars into the network output node list"""
  245. if not all([var.owner for var in vars]):
  246. self.add_dep_oprs(*vars)
  247. for var in vars:
  248. # use method 'is' instead of 'in' to avoid
  249. # compare VarNode use elemwise equal
  250. if not any(var is _ for _ in self.output_vars):
  251. self.output_vars.append(var)
  252. def remove_output(self, *vars: VarNode):
  253. r"""Removes vars from the network output node list"""
  254. for var in vars:
  255. # use list pop instead of remove to avoid
  256. # compare VarNode use elemwise equal
  257. is_removed = False
  258. for idx, out_var in enumerate(self.output_vars):
  259. if var is out_var:
  260. self.output_vars.pop(idx)
  261. is_removed = True
  262. if not is_removed:
  263. logger.warning(
  264. "Failed to remove {}({}). Please check whether "
  265. "this node is in the output list.".format(var.name, id(var))
  266. )
  267. def add_dep_oprs(self, *vars):
  268. if len(vars) == 0:
  269. vars = self.output_vars
  270. assert all(isinstance(var, VarNode) for var in vars), "Only support add VarNode"
  271. q = list(vars)
  272. while len(q) > 0:
  273. cur = q.pop(0)
  274. if cur.owner is not None:
  275. continue
  276. if cur.name is None:
  277. cur.name = cur.var.name
  278. self.all_vars_map[cur.var.id] = cur
  279. mge_opr = cur.var.owner
  280. if get_opr_type(mge_opr) == "Host2DeviceCopy":
  281. self._orig_inputs.extend(mge_opr.outputs)
  282. cur.owner = self._add_opr(mge_opr)
  283. if cur.owner is None:
  284. cur.owner = self.all_oprs_map[mge_opr.id]
  285. continue
  286. q.extend(cur.owner.inputs)
  287. return list(vars)
  288. def modify_opr_names(self, modifier):
  289. r"""Modifies names of operators **inplace**; useful for merging loaded
  290. network into another network
  291. Args:
  292. modifier(str or callable): a string to be prepended to the name, or a function
  293. that maps from name to name
  294. """
  295. if isinstance(modifier, str):
  296. om = modifier
  297. modifier = lambda v: "{}.{}".format(om, v)
  298. assert isinstance(modifier, collections.Callable)
  299. for i in self.all_oprs:
  300. v0 = i.name
  301. v1 = modifier(v0)
  302. assert isinstance(v1, str)
  303. i.name = v1
  304. def reset_batch_size(self, batchsize, *, blacklist=()):
  305. r"""Helper for reset batch size; first dimension of all data providers
  306. not in blacklist are assumed to be the batch size
  307. Args:
  308. blacklist: data provider names whose first dimension is not
  309. batchbatch size
  310. """
  311. blacklist = set(blacklist)
  312. prev_batchsize = None
  313. for i in self.data_providers_filter:
  314. if i.name in blacklist:
  315. blacklist.remove(i.name)
  316. else:
  317. shp = list(i.shape)
  318. if prev_batchsize is None:
  319. prev_batchsize = shp[0]
  320. else:
  321. assert prev_batchsize == shp[0], (
  322. "batchsize mismatch: batchsize={} "
  323. "shape={} dp={}".format(prev_batchsize, shp, i.name)
  324. )
  325. shp[0] = batchsize
  326. i.shape = tuple(shp)
  327. self._compile()
  328. assert prev_batchsize is not None, "no data provider found"
  329. assert not blacklist, "unused items in blacklist: {}".format(blacklist)
  330. def replace_vars(self, repl_dict: Dict[VarNode, VarNode]):
  331. r"""Replaces vars in the graph.
  332. Args:
  333. repl_dict: the map {old_var: new_var} that specifies how to replace the vars.
  334. """
  335. if not all([var.owner for var in repl_dict.values()]):
  336. self.add_dep_oprs(*list(repl_dict.values()))
  337. for var in self.all_vars:
  338. if var in repl_dict:
  339. repl_var = repl_dict[var]
  340. if repl_var is var:
  341. continue
  342. for opnode in var.users:
  343. # use method 'is' instead of 'in' to avoid
  344. # compare VarNode use elemwise equal
  345. assert any([var is _ for _ in opnode.inputs])
  346. opnode.inputs = [repl_var if var is i else i for i in opnode.inputs]
  347. if opnode not in repl_var.users:
  348. repl_var.users.append(opnode)
  349. var.users.clear()
  350. self._compile()
  351. def replace_oprs(self, repl_dict: Dict[OpNode, OpNode]):
  352. r"""Replaces operators in the graph.
  353. Args:
  354. repl_dict: the map {old_opr: new_opr} that specifies how to replace the operators.
  355. """
  356. for opr in self.all_oprs:
  357. if opr in repl_dict:
  358. assert len(opr.outputs) == len(
  359. repl_dict[opr].outputs
  360. ), "can not replace {} with {}".format(type(opr), type(repl_dict[opr]))
  361. for ind, var in enumerate(opr.outputs):
  362. var.owner = repl_dict[opr]
  363. var.__dict__.update(repl_dict[opr].outputs[ind].__dict__)
  364. var.var = repl_dict[opr].outputs[ind].var
  365. repl_dict[opr].outputs = opr.outputs
  366. self._compile()
  367. def get_opr_by_type(self, oprcls, unique=True):
  368. assert issubclass(oprcls, OpNode)
  369. rst = self.opr_filter.type(oprcls).as_list()
  370. if unique:
  371. assert len(rst) == 1, "{} operators of type {} found".format(
  372. len(rst), oprcls
  373. )
  374. (rst,) = rst
  375. return rst
  376. def get_opr_by_name(self, name, unique=True):
  377. rst = self.opr_filter.name(name).as_list()
  378. if unique:
  379. assert len(rst) == 1, "{} operators of type {} found".format(len(rst), name)
  380. (rst,) = rst
  381. return rst
  382. def get_var_by_name(self, name, unique=True):
  383. rst = self.var_filter.name(name).as_list()
  384. if unique:
  385. assert len(rst) == 1, "{} operators of type {} found".format(len(rst), name)
  386. (rst,) = rst
  387. return rst
  388. def get_var_receive_oprs(self, var):
  389. r"""Gets all oprs which use var as input"""
  390. return self.opr_filter.has_input(var).as_list()
  391. def get_dep_oprs(self, var):
  392. r"""Gets dependent oprs of var"""
  393. return get_oprs_seq(var, False, False)
  394. @property
  395. def opr_filter(self):
  396. r"""Filter on all opnodes of the Network."""
  397. oprs = self.all_oprs
  398. return NodeFilter(itertools.islice(oprs, len(oprs)))
  399. @property
  400. def var_filter(self):
  401. r"""Filter on all varnode of the Network."""
  402. vars = self.all_vars
  403. return NodeFilter(itertools.islice(vars, len(vars)))
  404. @property
  405. def params_filter(self): # all immutable tensor
  406. r"""Filter on all parameters (ImmutableTensor Opr) of the Network"""
  407. return self.opr_filter.param_provider()
  408. @property
  409. def data_providers_filter(self): # all host2devicecopy
  410. r"""Filter on all input nodes (Host2DeviceCopy Opr) of the Network"""
  411. return self.opr_filter.data_provider()
  412. @property
  413. def dest_vars(self):
  414. r"""Output varnodes of the Network."""
  415. return self.output_vars
  416. @property
  417. def all_oprs(self):
  418. return get_oprs_seq(self.output_vars, False, False)
  419. @property
  420. def all_vars(self):
  421. return get_dep_vars(self.output_vars)
  422. @property
  423. def all_vars_dict(self):
  424. return self.var_filter.as_dict()
  425. @property
  426. def all_oprs_dict(self):
  427. return self.opr_filter.as_dict()
  428. def _add_opr(self, opr) -> Optional[OpNode]:
  429. r"""Used for loading and building graph."""
  430. assert isinstance(opr, _imperative_rt.graph.OperatorNode)
  431. # TODO: use megbrain C++ RTTI to replace type string
  432. if opr.id not in self.all_oprs_map:
  433. opnode = str_to_mge_class(get_opr_type(opr)).load(opr)
  434. self.all_oprs_map[opr.id] = opnode
  435. for var in opr.inputs:
  436. varnode = self._get_var(var)
  437. opnode.add_inp_var(varnode)
  438. varnode.users.append(opnode)
  439. for var in opr.outputs:
  440. opnode.add_out_var(self._get_var(var))
  441. return opnode
  442. else:
  443. # overwrite the opnode 'new' output VarNode with
  444. # original one when output number larger than 1,
  445. # or will cause dependence issue in _compiler step.
  446. if len(opr.outputs) > 1:
  447. opnode = self.all_oprs_map[opr.id]
  448. for idx, output in enumerate(opnode.outputs):
  449. if output.var.id in self.all_vars_map:
  450. opnode.outputs[idx] = self.all_vars_map[output.var.id]
  451. return None
  452. def _get_opr(self, x):
  453. if x.id in self.all_oprs_map:
  454. return self.all_oprs_map[x.id]
  455. else:
  456. return None
  457. def _get_var(self, x):
  458. r"""Convert :class:`~._imperative_rt.graph.VarNode` to :class:`~.VarNode`."""
  459. assert isinstance(x, _imperative_rt.graph.VarNode)
  460. if x.id not in self.all_vars_map or self.all_vars_map[x.id].var != x:
  461. self.all_vars_map[x.id] = VarNode.load(x, self._get_opr(x.owner))
  462. return self.all_vars_map[x.id]
  463. def set_symbolic_shape(option: bool):
  464. r"""Set the VarNode use symbolic shape or not, return the last status.
  465. Please set to True and must recover after dump if want to change the input batch size.
  466. Args:
  467. option: True for enable symbolic shape.
  468. """
  469. return _set_symbolic_shape(option)
  470. def as_varnode(obj):
  471. r"""convert a :class:`.utils.network_node.VarNode` compatible object to :class:`.utils.network_node.VarNode`.
  472. Args:
  473. obj: it must be one of the following:
  474. 1. a :class:`.utils.network_node.VarNode` object
  475. 2. a :class:`.utils.network_node.OpNode` object that has unique output
  476. 3. an iterable that produces either type 1 or 2, with length 1
  477. """
  478. if type(obj) is VarNode:
  479. return obj
  480. if isinstance(obj, OpNode):
  481. assert len(obj.outputs) == 1, (
  482. "operator {} must have one output to be converted to VarNode; "
  483. "got {} actually".format(obj, len(obj.outputs))
  484. )
  485. ret = obj.outputs[0]
  486. assert type(ret) is VarNode
  487. return ret
  488. assert isinstance(
  489. obj, collections.Iterable
  490. ), "{} is not compatible with VarNode".format(obj)
  491. val = list(obj)
  492. assert (
  493. len(val) == 1
  494. ), "can not convert sequence of length {} to VarNode ({})".format(
  495. len(val), (lambda s: s if len(s) < 50 else s[:50] + " ...")(str(val))
  496. )
  497. return as_varnode(val[0])
  498. def as_oprnode(obj):
  499. r"""convert a :class:`.utils.network_node.OpNode` compatible object to
  500. :class:`.utils.network_node.OpNode`; it works like :func:`as_varnode`.
  501. """
  502. if type(obj) is VarNode:
  503. return obj.owner
  504. if isinstance(obj, OpNode):
  505. return obj
  506. assert isinstance(
  507. obj, collections.Iterable
  508. ), "{} is not compatible with OpNode".format(obj)
  509. val = list(obj)
  510. assert (
  511. len(val) == 1
  512. ), "can not convert sequence of length {} to " "OpNode({})".format(len(val), val)
  513. return as_oprnode(val[0])
  514. class NodeFilter:
  515. r"""Filter on node iterator. This class is an iterator of
  516. :class:`.NetworkNode` objects and multiple filtering conditions and
  517. mappers can be chained.
  518. Example:
  519. .. code-block::
  520. # find all :class:`.ImmutableTensor` nodes
  521. for i in NodeFilter(node_iter).param_provider():
  522. print(i)
  523. # find all :class:`.ImmutableTensor` nodes that end with ':W'
  524. for i in NodeFilter(node_iter).param_provider().name('*:W'):
  525. print(i)
  526. # number of inputs
  527. nr_input = NodeFilter(node_iter).data_provider().as_count()
  528. """
  529. _iter = None
  530. def __init__(self, node_iter):
  531. """
  532. :param node_iter: iterator to :class:`.NetworkNode`, or a
  533. :class:`.VarNode`-compatible object; in the later case, its
  534. dependent oprs would be used
  535. """
  536. if isinstance(node_iter, VarNode):
  537. oprs = get_oprs_seq(node_iter, False, False)
  538. node_iter = itertools.islice(oprs, len(oprs) - 1)
  539. if isinstance(node_iter, OpNode):
  540. oprs = get_oprs_seq(node_iter.inputs, False, False)
  541. node_iter = itertools.islice(oprs, len(oprs) - 1)
  542. assert isinstance(node_iter, collections.Iterable)
  543. if (not isinstance(node_iter, NodeFilter)) and type(
  544. self
  545. ) is not NodeFilterCheckType:
  546. node_iter = NodeFilterCheckType(node_iter, NetworkNode)
  547. self._iter = node_iter
  548. @classmethod
  549. def make_all_deps(cls, *dest_vars):
  550. r"""make a :class:`NodeFilter` that contains all deps of given vars"""
  551. return cls(list(get_oprs_seq(dest_vars, False, False)))
  552. def __iter__(self):
  553. r"""to be overwritten by subclass to implement filters"""
  554. return iter(self._iter)
  555. def type(self, node_type):
  556. r"""filter by specific node type
  557. Args:
  558. node_type: node type class
  559. Returns:
  560. a new :class:`NodeFilter` object
  561. """
  562. return NodeFilterType(self, node_type)
  563. def check_type(self, node_type):
  564. r"""assert that all oprs produced by this iterator are instances of
  565. certain type
  566. Args:
  567. node_type: node type class
  568. Returns:
  569. a new :class:`NodeFilter` object
  570. Raises:
  571. TypeError if type check failed
  572. """
  573. return NodeFilterCheckType(self, node_type)
  574. def not_type(self, node_type):
  575. r"""remove oprs of specific type
  576. Args:
  577. node_type: node type class
  578. Returns:
  579. a new :class:`NodeFilter` object
  580. """
  581. return NodeFilterNotType(self, node_type)
  582. def param_provider(self):
  583. r"""get :class:`~.ParamProvider` oprs; shorthand for
  584. ``.type(ParamProvider)``
  585. """
  586. return self.type(ImmutableTensor)
  587. def data_provider(self):
  588. r"""get :class:`.DataProvider` oprs; shorthand for
  589. ``.type(DataProvider)``
  590. """
  591. return self.type(Host2DeviceCopy)
  592. def name(self, pattern, ignorecase=True):
  593. r"""filter by node name
  594. Args:
  595. pattern(class:`str`): a string in glob syntax that can contain ``?`` and
  596. ``*`` to match a single or arbitrary characters.
  597. ignorecase(bool, optional): whether to ignroe case
  598. Returns:
  599. a new :class:`NodeFilter` object
  600. """
  601. return NodeFilterName(self, pattern, ignorecase)
  602. def has_input(self, var):
  603. r"""an opr is kept if it has given var as one of its inputs
  604. Args:
  605. var: var node to checked
  606. Returns:
  607. a new :class:`NodeFilter` object
  608. """
  609. return NodeFilterHasInput(self, var)
  610. def as_list(self):
  611. r"""consume this iterator and return its content as a list"""
  612. return list(self)
  613. def as_unique(self):
  614. r"""assert that this iterator yields only one node and return it
  615. Returns:
  616. class:`.GraphNodeBase`: the unique node
  617. Raises:
  618. ValueError if this iterator does not yield a unique node
  619. """
  620. (opr,) = self
  621. return opr
  622. def as_dict(self):
  623. r"""construct an ordered dict to map from node names to objects in
  624. this iterator
  625. """
  626. return collections.OrderedDict((i.name, i) for i in self)
  627. def as_count(self):
  628. r"""consume this iterator and get the number of elements"""
  629. return sum(1 for _ in self)
  630. class NodeFilterType(NodeFilter):
  631. r"""see :meth:`NodeFilter.type`"""
  632. _node_type = None
  633. def __init__(self, node_iter, node_type):
  634. assert issubclass(node_type, NetworkNode), "bad opr type: {}".format(node_type)
  635. super().__init__(node_iter)
  636. self._node_type = node_type
  637. def __iter__(self):
  638. for i in self._iter:
  639. if isinstance(i, self._node_type):
  640. yield i
  641. class NodeFilterNotType(NodeFilterType):
  642. r"""see :meth:`NodeFilter.not_type`"""
  643. def __iter__(self):
  644. for i in self._iter:
  645. if not isinstance(i, self._node_type):
  646. yield i
  647. class NodeFilterCheckType(NodeFilterType):
  648. r"""see :meth:`NodeFilter.check_type`"""
  649. def __iter__(self):
  650. for i in self._iter:
  651. if not isinstance(i, self._node_type):
  652. raise TypeError(
  653. "all nodes should be {}; got {!r}".format(self._node_type, i)
  654. )
  655. yield i
  656. class NodeFilterHasInput(NodeFilter):
  657. r"""see :meth:`NodeFilter.has_input`"""
  658. _var = None
  659. def __init__(self, node_iter, var):
  660. var = as_varnode(var)
  661. super().__init__(node_iter)
  662. self.var = var
  663. def __iter__(self):
  664. for i in self._iter:
  665. assert isinstance(
  666. i, OpNode
  667. ), "has_input() must be used with OpNode; " "got {!r}".format(i)
  668. if any(self.var is _ for _ in i.inputs):
  669. yield i
  670. class NodeFilterName(NodeFilter):
  671. r"""see :meth:`NodeFilter.name`"""
  672. _re = None
  673. def __init__(self, node_iter, pattern, ignorecase):
  674. super().__init__(node_iter)
  675. self.pattern = pattern
  676. self._re = self.make_re(pattern, ignorecase)
  677. @classmethod
  678. def make_re(cls, pattern, ignorecase=True):
  679. assert isinstance(pattern, str), "bad pattern: {!r}".format(pattern)
  680. assert isinstance(ignorecase, bool)
  681. flags = 0
  682. if ignorecase:
  683. flags |= re.IGNORECASE
  684. return re.compile(fnmatch.translate(pattern), flags=flags)
  685. def __iter__(self):
  686. for i in self._iter:
  687. if self.pattern == i.name or self._re.match(i.name):
  688. yield i

MegEngine 安装包中集成了使用 GPU 运行代码所需的 CUDA 环境,不用区分 CPU 和 GPU 版。 如果想要运行 GPU 程序,请确保机器本身配有 GPU 硬件设备并安装好驱动。 如果你想体验在云端 GPU 算力平台进行深度学习开发的感觉,欢迎访问 MegStudio 平台