| ## @package net_printer | |
| # Module caffe2.python.net_printer | |
| from caffe2.proto.caffe2_pb2 import OperatorDef, NetDef | |
| from caffe2.python.checkpoint import Job | |
| from caffe2.python.core import Net, ExecutionStep, Plan | |
| from caffe2.python.task import Task, TaskGroup, WorkspaceType, TaskOutput | |
| from collections import defaultdict | |
| from contextlib import contextmanager | |
| from copy import copy | |
| from itertools import chain | |
| class Visitor: | |
| @classmethod | |
| def register(cls, Type): | |
| if not(hasattr(cls, 'visitors')): | |
| cls.visitors = {} | |
| else: | |
| assert Type not in cls.visitors, \ | |
| '{} already registered!'.format(Type) | |
| def _register(func): | |
| cls.visitors[Type] = func | |
| return func | |
| return _register | |
| def __call__(self, obj, *args, **kwargs): | |
| if obj is None: | |
| return | |
| Type = type(obj) | |
| if Type not in self.__class__.visitors: | |
| raise TypeError('%s: unsupported object type: %s' % ( | |
| self.__class__.__name__, Type)) | |
| func = self.__class__.visitors[Type] | |
| return func(self, obj, *args, **kwargs) | |
| class Analyzer(Visitor): | |
| PREFIXES_TO_IGNORE = {'distributed_ctx_init'} | |
| def __init__(self): | |
| self.workspaces = defaultdict(lambda: defaultdict(lambda: 0)) | |
| self.workspace_ctx = [] | |
| @property | |
| def workspace(self): | |
| return self.workspace_ctx[-1] | |
| @contextmanager | |
| def set_workspace(self, node=None, ws=None, do_copy=False): | |
| if ws is not None: | |
| ws = ws | |
| elif node is not None: | |
| ws = self.workspaces[str(node)] | |
| else: | |
| ws = self.workspace | |
| if do_copy: | |
| ws = copy(ws) | |
| self.workspace_ctx.append(ws) | |
| try: | |
| yield ws | |
| finally: | |
| del self.workspace_ctx[-1] | |
| def define_blob(self, blob): | |
| self.workspace[blob] += 1 | |
| def need_blob(self, blob): | |
| if any(blob.startswith(p) for p in Analyzer.PREFIXES_TO_IGNORE): | |
| return | |
| assert blob in self.workspace, 'Blob undefined: %s' % blob | |
| @Analyzer.register(OperatorDef) | |
| def analyze_op(analyzer, op): | |
| for x in op.input: | |
| analyzer.need_blob(x) | |
| for x in op.output: | |
| analyzer.define_blob(x) | |
| @Analyzer.register(Net) | |
| def analyze_net(analyzer, net): | |
| for x in net.Proto().op: | |
| analyzer(x) | |
| @Analyzer.register(ExecutionStep) | |
| def analyze_step(analyzer, step): | |
| proto = step.Proto() | |
| with analyzer.set_workspace(do_copy=proto.create_workspace): | |
| if proto.report_net: | |
| with analyzer.set_workspace(do_copy=True): | |
| analyzer(step.get_net(proto.report_net)) | |
| all_new_blobs = set() | |
| substeps = step.Substeps() + [step.get_net(n) for n in proto.network] | |
| for substep in substeps: | |
| with analyzer.set_workspace( | |
| do_copy=proto.concurrent_substeps) as ws_in: | |
| analyzer(substep) | |
| if proto.should_stop_blob: | |
| analyzer.need_blob(proto.should_stop_blob) | |
| if proto.concurrent_substeps: | |
| new_blobs = set(ws_in.keys()) - set(analyzer.workspace.keys()) | |
| assert len(all_new_blobs & new_blobs) == 0, ( | |
| 'Error: Blobs created by multiple parallel steps: %s' % ( | |
| ', '.join(all_new_blobs & new_blobs))) | |
| all_new_blobs |= new_blobs | |
| for x in all_new_blobs: | |
| analyzer.define_blob(x) | |
| @Analyzer.register(Task) | |
| def analyze_task(analyzer, task): | |
| # check that our plan protobuf is not too large (limit of 64Mb) | |
| step = task.get_step() | |
| plan = Plan(task.node) | |
| plan.AddStep(step) | |
| proto_len = len(plan.Proto().SerializeToString()) | |
| assert proto_len < 2 ** 26, ( | |
| 'Due to a protobuf limitation, serialized tasks must be smaller ' | |
| 'than 64Mb, but this task has {} bytes.' % proto_len) | |
| is_private = task.workspace_type() != WorkspaceType.GLOBAL | |
| with analyzer.set_workspace(do_copy=is_private): | |
| analyzer(step) | |
| @Analyzer.register(TaskGroup) | |
| def analyze_task_group(analyzer, tg): | |
| for task in tg.tasks_by_node().tasks(): | |
| with analyzer.set_workspace(node=task.node): | |
| analyzer(task) | |
| @Analyzer.register(Job) | |
| def analyze_job(analyzer, job): | |
| analyzer(job.init_group) | |
| analyzer(job.epoch_group) | |
| def analyze(obj): | |
| """ | |
| Given a Job, visits all the execution steps making sure that: | |
| - no undefined blobs will be found during execution | |
| - no blob with same name is defined in concurrent steps | |
| """ | |
| Analyzer()(obj) | |
| class Text: | |
| def __init__(self): | |
| self._indent = 0 | |
| self._lines_in_context = [0] | |
| self.lines = [] | |
| @contextmanager | |
| def context(self, text): | |
| if text is not None: | |
| self.add('with %s:' % text) | |
| self._indent += 4 | |
| self._lines_in_context.append(0) | |
| try: | |
| yield | |
| finally: | |
| if text is not None: | |
| if self._lines_in_context[-1] == 0: | |
| self.add('pass') | |
| self._indent -= 4 | |
| del self._lines_in_context[-1] | |
| def add(self, text): | |
| self._lines_in_context[-1] += 1 | |
| self.lines.append((' ' * self._indent) + text) | |
| def __str__(self): | |
| return '\n'.join(self.lines) | |
| class Printer(Visitor, Text): | |
| def __init__(self, factor_prefixes=False, c2_syntax=True): | |
| super(Visitor, self).__init__() | |
| super(Text, self).__init__() | |
| self.factor_prefixes = factor_prefixes | |
| self.c2_syntax = c2_syntax | |
| self.c2_net_name = None | |
| def _sanitize_str(s): | |
| if isinstance(s, str): | |
| sanitized = s | |
| elif isinstance(s, bytes): | |
| sanitized = s.decode('ascii', errors='ignore') | |
| else: | |
| sanitized = str(s) | |
| if len(sanitized) < 64: | |
| return "'%s'" % sanitized | |
| else: | |
| return "'%s'" % sanitized[:64] + '...<+len=%d>' % (len(sanitized) - 64) | |
| def _arg_val(arg): | |
| if arg.HasField('f'): | |
| return str(arg.f) | |
| if arg.HasField('i'): | |
| return str(arg.i) | |
| if arg.HasField('s'): | |
| return _sanitize_str(arg.s) | |
| if arg.floats: | |
| return str(list(arg.floats)) | |
| if arg.ints: | |
| return str(list(arg.ints)) | |
| if arg.strings: | |
| return str([_sanitize_str(s) for s in arg.strings]) | |
| return '[]' | |
| def commonprefix(m): | |
| "Given a list of strings, returns the longest common prefix" | |
| if not m: | |
| return '' | |
| s1 = min(m) | |
| s2 = max(m) | |
| for i, c in enumerate(s1): | |
| if c != s2[i]: | |
| return s1[:i] | |
| return s1 | |
| def format_value(val): | |
| if isinstance(val, list): | |
| return '[%s]' % ', '.join("'%s'" % str(v) for v in val) | |
| else: | |
| return str(val) | |
| def factor_prefix(vals, do_it): | |
| vals = [format_value(v) for v in vals] | |
| prefix = commonprefix(vals) if len(vals) > 1 and do_it else '' | |
| joined = ', '.join(v[len(prefix):] for v in vals) | |
| return '%s[%s]' % (prefix, joined) if prefix else joined | |
| def call(op, inputs=None, outputs=None, factor_prefixes=False): | |
| if not inputs: | |
| inputs = '' | |
| else: | |
| inputs_v = [a for a in inputs if not isinstance(a, tuple)] | |
| inputs_kv = [a for a in inputs if isinstance(a, tuple)] | |
| inputs = ', '.join( | |
| x | |
| for x in chain( | |
| [factor_prefix(inputs_v, factor_prefixes)], | |
| ('%s=%s' % kv for kv in inputs_kv), | |
| ) | |
| if x | |
| ) | |
| call = '%s(%s)' % (op, inputs) | |
| return call if not outputs else '%s = %s' % ( | |
| factor_prefix(outputs, factor_prefixes), call) | |
| def format_device_option(dev_opt): | |
| if not dev_opt or not ( | |
| dev_opt.device_type or dev_opt.device_id or dev_opt.node_name): | |
| return None | |
| return call( | |
| 'DeviceOption', | |
| [dev_opt.device_type, dev_opt.device_id, "'%s'" % dev_opt.node_name]) | |
| @Printer.register(OperatorDef) | |
| def print_op(text, op): | |
| args = [(a.name, _arg_val(a)) for a in op.arg] | |
| dev_opt_txt = format_device_option(op.device_option) | |
| if dev_opt_txt: | |
| args.append(('device_option', dev_opt_txt)) | |
| if text.c2_net_name: | |
| text.add(call( | |
| text.c2_net_name + '.' + op.type, | |
| [list(op.input), list(op.output)] + args)) | |
| else: | |
| text.add(call( | |
| op.type, | |
| list(op.input) + args, | |
| op.output, | |
| factor_prefixes=text.factor_prefixes)) | |
| for arg in op.arg: | |
| if arg.HasField('n'): | |
| with text.context('arg: %s' % arg.name): | |
| text(arg.n) | |
| @Printer.register(NetDef) | |
| def print_net_def(text, net_def): | |
| if text.c2_syntax: | |
| text.add(call('core.Net', ["'%s'" % net_def.name], [net_def.name])) | |
| text.c2_net_name = net_def.name | |
| else: | |
| text.add('# net: %s' % net_def.name) | |
| for op in net_def.op: | |
| text(op) | |
| if text.c2_syntax: | |
| text.c2_net_name = None | |
| @Printer.register(Net) | |
| def print_net(text, net): | |
| text(net.Proto()) | |
| def _get_step_context(step): | |
| proto = step.Proto() | |
| if proto.should_stop_blob: | |
| return call('loop'), False | |
| if proto.num_iter and proto.num_iter != 1: | |
| return call('loop', [proto.num_iter]), False | |
| if proto.num_concurrent_instances > 1: | |
| return ( | |
| call('parallel', | |
| [('num_instances', proto.num_concurrent_instances)]), | |
| len(step.Substeps()) > 1) | |
| concurrent = proto.concurrent_substeps and len(step.Substeps()) > 1 | |
| if concurrent: | |
| return call('parallel'), True | |
| if proto.report_net: | |
| return call('run_once'), False | |
| return None, False | |
| @Printer.register(ExecutionStep) | |
| def print_step(text, step): | |
| proto = step.Proto() | |
| step_ctx, do_substep = _get_step_context(step) | |
| with text.context(step_ctx): | |
| if proto.report_net: | |
| with text.context(call('report_net', [proto.report_interval])): | |
| text(step.get_net(proto.report_net)) | |
| substeps = step.Substeps() + [step.get_net(n) for n in proto.network] | |
| for substep in substeps: | |
| sub_proto = ( | |
| substep.Proto() if isinstance(substep, ExecutionStep) else None) | |
| if sub_proto is not None and sub_proto.run_every_ms: | |
| substep_ctx = call( | |
| 'reporter', | |
| [str(substep), ('interval_ms', sub_proto.run_every_ms)]) | |
| elif do_substep: | |
| title = ( | |
| 'workspace' | |
| if sub_proto is not None and sub_proto.create_workspace else | |
| 'step') | |
| substep_ctx = call(title, [str(substep)]) | |
| else: | |
| substep_ctx = None | |
| with text.context(substep_ctx): | |
| text(substep) | |
| if proto.should_stop_blob: | |
| text.add(call('yield stop_if', [proto.should_stop_blob])) | |
| def _print_task_output(x): | |
| assert isinstance(x, TaskOutput) | |
| return 'Output[' + ', '.join(str(x) for x in x.names) + ']' | |
| @Printer.register(Task) | |
| def print_task(text, task): | |
| outs = ', '.join(_print_task_output(o) for o in task.outputs()) | |
| context = [('node', task.node), ('name', task.name), ('outputs', outs)] | |
| with text.context(call('Task', context)): | |
| text(task.get_step()) | |
| @Printer.register(TaskGroup) | |
| def print_task_group(text, tg, header=None): | |
| with text.context(header or call('TaskGroup')): | |
| for task in tg.tasks_by_node().tasks(): | |
| text(task) | |
| @Printer.register(Job) | |
| def print_job(text, job): | |
| text(job.init_group, 'Job.current().init_group') | |
| text(job.epoch_group, 'Job.current().epoch_group') | |
| with text.context('Job.current().stop_conditions'): | |
| for out in job.stop_conditions: | |
| text.add(_print_task_output(out)) | |
| text(job.download_group, 'Job.current().download_group') | |
| text(job.exit_group, 'Job.current().exit_group') | |
| def to_string(obj, **kwargs): | |
| """ | |
| Given a Net, ExecutionStep, Task, TaskGroup or Job, produces a string | |
| with detailed description of the execution steps. | |
| """ | |
| printer = Printer(**kwargs) | |
| printer(obj) | |
| return str(printer) | |
| def debug_net(net): | |
| """ | |
| Given a Net, produce another net that logs info about the operator call | |
| before each operator execution. Use for debugging purposes. | |
| """ | |
| assert isinstance(net, Net) | |
| debug_net = Net(str(net)) | |
| assert isinstance(net, Net) | |
| for op in net.Proto().op: | |
| text = Text() | |
| print_op(op, text) | |
| debug_net.LogInfo(str(text)) | |
| debug_net.Proto().op.extend([op]) | |
| return debug_net |