Caffe-HRT/python/caffe/draw.py at master · bjthemost/Caffe-HRT · GitHub

Repository navigation

History

244 lines (213 loc) · 8.76 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

"""

Caffe network visualization: draw the NetParameter protobuffer.

.. note::

This requires pydot>=1.0.2, which is not included in requirements.txt since

it requires graphviz and other prerequisites outside the scope of the

Caffe.

"""

from caffe.proto import caffe_pb2

"""

pydot is not supported under python 3 and pydot2 doesn't work properly.

pydotplus works nicely (pip install pydotplus)

"""

try:

# Try to load pydotplus

import pydotplus as pydot

except ImportError:

import pydot

# Internal layer and blob styles.

LAYER_STYLE_DEFAULT = {'shape': 'record',

'fillcolor': '#6495ED',

'style': 'filled'}

NEURON_LAYER_STYLE = {'shape': 'record',

'fillcolor': '#90EE90',

'style': 'filled'}

BLOB_STYLE = {'shape': 'octagon',

'fillcolor': '#E0E0E0',

'style': 'filled'}

def get_pooling_types_dict():

"""Get dictionary mapping pooling type number to type name

"""

desc = caffe_pb2.PoolingParameter.PoolMethod.DESCRIPTOR

d = {}

for k, v in desc.values_by_name.items():

d[v.number] = k

return d

def get_edge_label(layer):

"""Define edge label based on layer type.

"""

if layer.type == 'Data':

edge_label = 'Batch ' + str(layer.data_param.batch_size)

elif layer.type == 'Convolution' or layer.type == 'Deconvolution':

edge_label = str(layer.convolution_param.num_output)

elif layer.type == 'InnerProduct':

edge_label = str(layer.inner_product_param.num_output)

else:

edge_label = '""'

return edge_label

def get_layer_label(layer, rankdir):

"""Define node label based on layer type.

Parameters

----------

layer : ?

rankdir : {'LR', 'TB', 'BT'}

Direction of graph layout.

Returns

-------

string :

A label for the current layer

"""

if rankdir in ('TB', 'BT'):

# If graph orientation is vertical, horizontal space is free and

# vertical space is not; separate words with spaces

separator = ' '

else:

# If graph orientation is horizontal, vertical space is free and

# horizontal space is not; separate words with newlines

separator = '\\n'

if layer.type == 'Convolution' or layer.type == 'Deconvolution':

# Outer double quotes needed or else colon characters don't parse

# properly

node_label = '"%s%s(%s)%skernel size: %d%sstride: %d%spad: %d"' %\

(layer.name,

separator,

layer.type,

separator,

layer.convolution_param.kernel_size[0] if len(layer.convolution_param.kernel_size._values) else 1,

separator,

layer.convolution_param.stride[0] if len(layer.convolution_param.stride._values) else 1,

separator,

layer.convolution_param.pad[0] if len(layer.convolution_param.pad._values) else 0)

elif layer.type == 'Pooling':

pooling_types_dict = get_pooling_types_dict()

node_label = '"%s%s(%s %s)%skernel size: %d%sstride: %d%spad: %d"' %\

(layer.name,

separator,

pooling_types_dict[layer.pooling_param.pool],

layer.type,

separator,

layer.pooling_param.kernel_size[0] if len(layer.pooling_param.kernel_size._values) else 1,

separator,

layer.pooling_param.stride[0] if len(layer.pooling_param.stride._values) else 1,

separator,

layer.pooling_param.pad[0] if len(layer.pooling_param.pad._values) else 0)

else:

node_label = '"%s%s(%s)"' % (layer.name, separator, layer.type)

return node_label

def choose_color_by_layertype(layertype):

"""Define colors for nodes based on the layer type.

"""

color = '#6495ED' # Default

if layertype == 'Convolution' or layertype == 'Deconvolution':

color = '#FF5050'

elif layertype == 'Pooling':

color = '#FF9900'

elif layertype == 'InnerProduct':

color = '#CC33FF'

return color

def get_pydot_graph(caffe_net, rankdir, label_edges=True, phase=None):

"""Create a data structure which represents the `caffe_net`.

Parameters

----------

caffe_net : object

rankdir : {'LR', 'TB', 'BT'}

Direction of graph layout.

label_edges : boolean, optional

Label the edges (default is True).

phase : {caffe_pb2.Phase.TRAIN, caffe_pb2.Phase.TEST, None} optional

Include layers from this network phase. If None, include all layers.

(the default is None)

Returns

-------

pydot graph object

"""

pydot_graph = pydot.Dot(caffe_net.name if caffe_net.name else 'Net',

graph_type='digraph',

rankdir=rankdir)

pydot_nodes = {}

pydot_edges = []

for layer in caffe_net.layer:

if phase is not None:

included = False

if len(layer.include) == 0:

included = True

if len(layer.include) > 0 and len(layer.exclude) > 0:

raise ValueError('layer ' + layer.name + ' has both include '

'and exclude specified.')

for layer_phase in layer.include:

included = included or layer_phase.phase == phase

for layer_phase in layer.exclude:

included = included and not layer_phase.phase == phase

if not included:

continue

node_label = get_layer_label(layer, rankdir)

node_name = "%s_%s" % (layer.name, layer.type)

if (len(layer.bottom) == 1 and len(layer.top) == 1 and

layer.bottom[0] == layer.top[0]):

# We have an in-place neuron layer.

pydot_nodes[node_name] = pydot.Node(node_label,

**NEURON_LAYER_STYLE)

else:

layer_style = LAYER_STYLE_DEFAULT

layer_style['fillcolor'] = choose_color_by_layertype(layer.type)

pydot_nodes[node_name] = pydot.Node(node_label, **layer_style)

for bottom_blob in layer.bottom:

pydot_nodes[bottom_blob + '_blob'] = pydot.Node('%s' % bottom_blob,

**BLOB_STYLE)

edge_label = '""'

pydot_edges.append({'src': bottom_blob + '_blob',

'dst': node_name,

'label': edge_label})

for top_blob in layer.top:

pydot_nodes[top_blob + '_blob'] = pydot.Node('%s' % (top_blob))

if label_edges:

edge_label = get_edge_label(layer)

else:

edge_label = '""'

pydot_edges.append({'src': node_name,

'dst': top_blob + '_blob',

'label': edge_label})

# Now, add the nodes and edges to the graph.

for node in pydot_nodes.values():

pydot_graph.add_node(node)

for edge in pydot_edges:

pydot_graph.add_edge(

pydot.Edge(pydot_nodes[edge['src']],

pydot_nodes[edge['dst']],

label=edge['label']))

return pydot_graph

def draw_net(caffe_net, rankdir, ext='png', phase=None):

"""Draws a caffe net and returns the image string encoded using the given

extension.

Parameters

----------

caffe_net : a caffe.proto.caffe_pb2.NetParameter protocol buffer.

ext : string, optional

The image extension (the default is 'png').

phase : {caffe_pb2.Phase.TRAIN, caffe_pb2.Phase.TEST, None} optional

Include layers from this network phase. If None, include all layers.

(the default is None)

Returns

-------

string :

Postscript representation of the graph.

"""

return get_pydot_graph(caffe_net, rankdir, phase=phase).create(format=ext)

def draw_net_to_file(caffe_net, filename, rankdir='LR', phase=None):

"""Draws a caffe net, and saves it to file using the format given as the

file extension. Use '.raw' to output raw text that you can manually feed

to graphviz to draw graphs.

Parameters

----------

caffe_net : a caffe.proto.caffe_pb2.NetParameter protocol buffer.

filename : string

The path to a file where the networks visualization will be stored.

rankdir : {'LR', 'TB', 'BT'}

Direction of graph layout.

phase : {caffe_pb2.Phase.TRAIN, caffe_pb2.Phase.TEST, None} optional

Include layers from this network phase. If None, include all layers.

(the default is None)

"""

ext = filename[filename.rfind('.')+1:]

with open(filename, 'wb') as fid:

fid.write(draw_net(caffe_net, rankdir, ext, phase))

Caffe-HRT/python/caffe/draw.py at master · bjthemost/Caffe-HRT · GitHub

Repository navigation

Expand file tree

Breadcrumbs

draw.py

Latest commit

History

Breadcrumbs

draw.py

File metadata and controls