Caffe2 - Python API
A deep learning, cross platform ML framework
generator.py
1 ## @package generator
2 # Module caffe2.python.docs.generator
3 from __future__ import absolute_import
4 from __future__ import division
5 from __future__ import print_function
6 from __future__ import unicode_literals
7 import argparse
8 import os
9 from caffe2.python import core, workspace
10 from caffe2.python.docs.formatter import Markdown
11 from future.utils import viewitems, viewvalues
12 
13 OpSchema = workspace.C.OpSchema
14 
15 
16 class DocUploader(object):
17  def __init__(self):
18  pass
19 
20  def upload(self, text):
21  pass
22 
23 
24 class DocGenerator(object):
25  def __init__(self, formatter, uploader):
26  self.formatter = formatter
27  self.uploader = uploader
28  self.content_body = ""
29 
30  def create_body(self):
31  pass
32 
33  def update(self):
34  self.uploader.upload(self.content_body)
35 
36 
38  def getOperatorDoc(self, name, schema, priority):
39  return OperatorDoc(name, schema, priority)
40 
41  def getOperatorEngine(self, name):
42  return OperatorEngine(name)
43 
44  def getOperators(self):
45  # map: op_name -> operator
46  self.operators = {}
47  # map: op_name -> [engine, engine]
48  self.engines = {}
49 
50  def filePriority(x):
51  if x == "pytorch/pytorch/operators":
52  return 0
53  if 'contrib' in x.split('/'):
54  return 2
55  if 'experiments' in x.split('/'):
56  return 3
57  return 1
58 
59  for name in core._GetRegisteredOperators():
60  schema = OpSchema.get(name)
61  if schema:
62  priority = filePriority(os.path.dirname(schema.file))
63  operator = self.getOperatorDoc(name, schema, priority)
64  self.operators[name] = operator
65 
66  # Engine
67  elif name.find("_ENGINE_") != -1:
68  engine = self.getOperatorEngine(name)
69  if engine.base_op_name in self.engines:
70  self.engines[engine.base_op_name].append(engine)
71  else:
72  self.engines[engine.base_op_name] = [engine]
73 
74  # No schema
75  else:
76  priority = 4
77  self.operators[name] = self.getOperatorDoc(name, schema, priority)
78 
79  for name, engines in viewitems(self.engines):
80  if name in self.operators:
81  self.operators[name].addEngines(engines)
82 
83  # Generate a sorted list of operators
84  return sorted(
85  viewvalues(self.operators),
86  key=lambda op: (op.priority, op.name)
87  )
88 
89  def createBody(self):
90  operators = self.getOperators()
91 
92  for operator in operators:
93  operator.generateSchema(self.formatter)
94 
95  self.content_body += self.formatter.dump()
96 
97 
98 class OperatorEngine(object):
99  def __init__(self, name):
100  self.op_name = name
101  self.base_op_name, self.engine = name.split("_ENGINE_", 1)
102 
103  def getDeviceImpl(self):
104  deviceImplList = []
105  for device, impl in [('CPU', OpSchema.get_cpu_impl(self.op_name)),
106  ('CUDA', OpSchema.get_cuda_impl(self.op_name))]:
107  if not impl:
108  continue
109  deviceImplList.append((device, impl))
110  return deviceImplList
111 
112  def generateDoc(self, formatter):
113  for device, impl in self.getDeviceImpl():
114  formatter.addLine(
115  '{engine} on {device}: {impl}'.format(engine=self.engine,
116  device=device,
117  impl=impl))
118 
119 
120 class OperatorDoc(object):
121  def __init__(self, name, schema, priority):
122  self.name = name
123  self.schema = schema
124  self.priority = priority
125  print("Gathering docs for {}...".format(self.name))
126  self.engines = []
127 
128  def addEngines(self, engines):
129  self.engines = engines
130 
131  def generateDoc(self, formatter):
132  if self.schema.doc:
133  formatter.parseAndAdd(self.schema.doc)
134  formatter.addLinebreak()
135  else:
136  formatter.addLine("No documentation yet.")
137 
138  def generateTable(self, formatter, tuples, title_row, title):
139  if tuples:
140  if title:
141  formatter.addHeader(title, 3)
142  table = []
143  if title_row:
144  table = [title_row]
145  for name, doc in tuples:
146  table.append([name, doc or ''])
147  formatter.addTable(table, (table == []))
148 
149  def generateInterface(self, formatter):
150  def makeDesc(title, args):
151  f = formatter.clone()
152  f.addEmphasis(title, 1)
153  out = [(f.dump(), '')]
154  for arg in args:
155  f = formatter.clone()
156  if isinstance(arg, tuple):
157  name = arg[0]
158  if len(arg) > 1:
159  description = arg[1] or ''
160  else:
161  description = ''
162  else:
163  name = arg.name
164  description = arg.description or ''
165  f.addCode(name, inline=True)
166  out.append((f.dump(), description or ''))
167  return out
168 
169  tuples = []
170 
171  if self.schema.args:
172  tuples += makeDesc('Arguments', self.schema.args)
173 
174  if self.schema.input_desc:
175  tuples += makeDesc('Inputs', self.schema.input_desc)
176 
177  if self.schema.output_desc:
178  tuples += makeDesc('Outputs', self.schema.output_desc)
179 
180  self.generateTable(formatter, tuples, None, 'Interface')
181  print("Generated interface for {}".format(self.name))
182 
183  def generateCodeLink(self, formatter):
184  formatter.addHeader("Code", 3)
185  formatter.addLinebreak()
186  formatter.addCodeLink(self.schema.file)
187 
188  def getInfo(self, formatter, name, impl):
189  pass
190 
191  def generateDevices(self, formatter):
192  formatter.addHeader("Devices", 3)
193  devices = [
194  self.getInfo(formatter,
195  'CPU', OpSchema.get_cpu_impl(self.name)),
196  self.getInfo(formatter,
197  'GPU', OpSchema.get_cuda_impl(self.name)),
198  ]
199  formatter.addList([i for i in devices if i])
200 
201  def generateEngines(self, formatter):
202  if not len(self.engines):
203  return
204  formatter.addHeader("Engines", 3)
205  for engine in self.engines:
206  engine.generateDoc(formatter)
207 
208  def generateSchema(self, formatter):
209  formatter.addHeader(self.name, 2)
210  if self.schema:
211  self.generateDoc(formatter)
212  self.generateInterface(formatter)
213  self.generateCodeLink(formatter)
214  self.generateDevices(formatter)
215  self.generateEngines(formatter)
216  formatter.addBreak()
217  else:
218  formatter.addLine("No schema documented yet.")
219  self.generateDevices(formatter)
220 
221 
222 if __name__ == "__main__":
223  parser = argparse.ArgumentParser(description="Operators catalog generator.")
224  parser.add_argument('catalog_path', type=str,
225  help='operators-catalogue.md to write out to')
226  args = parser.parse_args()
227 
228  with open(args.catalog_path, 'w') as fp:
229  ops = OpDocGenerator(Markdown(), DocUploader())
230  ops.createBody()
231  fp.write(ops.content_body)
def getOperatorDoc(self, name, schema, priority)
Definition: generator.py:38
def getOperatorEngine(self, name)
Definition: generator.py:41
def generateTable(self, formatter, tuples, title_row, title)
Definition: generator.py:138