Skip to content

Commit 4ba40d7

Browse files
author
Saumitro Dasgupta
committed
Split core into graph and data + replaced indexing constants with named tuples
1 parent c0100fc commit 4ba40d7

File tree

8 files changed

+167
-179
lines changed

8 files changed

+167
-179
lines changed

‎kaffe/__init__.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1-
from .base import KaffeError, print_stderr
2-
from .core import GraphBuilder, DataReshaper, NodeMapper
1+
from .graph import GraphBuilder, NodeMapper
2+
from .data import DataReshaper
3+
from .errors import KaffeError, print_stderr
4+
35
from . import tensorflow

‎kaffe/base.py

Lines changed: 0 additions & 21 deletions
This file was deleted.

‎kaffe/data.py

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import numpy as np
2+
3+
from .caffe import get_caffe_resolver, has_pycaffe
4+
from .errors import KaffeError, print_stderr
5+
from .layers import NodeKind
6+
7+
class DataInjector(object):
8+
9+
def __init__(self, def_path, data_path):
10+
self.def_path = def_path
11+
self.data_path = data_path
12+
self.did_use_pb = False
13+
self.params = None
14+
self.load()
15+
16+
def load(self):
17+
if has_pycaffe():
18+
self.load_using_caffe()
19+
else:
20+
self.load_using_pb()
21+
22+
def load_using_caffe(self):
23+
caffe = get_caffe_resolver().caffe
24+
net = caffe.Net(self.def_path, self.data_path, caffe.TEST)
25+
data = lambda blob: blob.data
26+
self.params = [(k, map(data, v)) for k, v in net.params.items()]
27+
28+
def load_using_pb(self):
29+
data = get_caffe_resolver().NetParameter()
30+
data.MergeFromString(open(self.data_path, 'rb').read())
31+
pair = lambda layer: (layer.name, self.transform_data(layer))
32+
layers = data.layers or data.layer
33+
self.params = [pair(layer) for layer in layers if layer.blobs]
34+
self.did_use_pb = True
35+
36+
def transform_data(self, layer):
37+
transformed = []
38+
for blob in layer.blobs:
39+
if len(blob.shape.dim):
40+
dims = blob.shape.dim
41+
c_o, c_i, h, w = map(int, [1] * (4 - len(dims)) + list(dims))
42+
else:
43+
c_o = blob.num
44+
c_i = blob.channels
45+
h = blob.height
46+
w = blob.width
47+
data = np.array(blob.data, dtype=np.float32).reshape(c_o, c_i, h, w)
48+
transformed.append(data)
49+
return transformed
50+
51+
def adjust_parameters(self, node, data):
52+
if not self.did_use_pb:
53+
return data
54+
# When using the protobuf-backend, each parameter initially has four dimensions.
55+
# In certain cases (like FC layers), we want to eliminate the singleton dimensions.
56+
# This implementation takes care of the common cases. However, it does leave the
57+
# potential for future issues.
58+
# The Caffe-backend does not suffer from this problem.
59+
data = list(data)
60+
squeeze_indices = [1] # Squeeze biases.
61+
if node.kind == NodeKind.InnerProduct:
62+
squeeze_indices.append(0) # Squeeze FC.
63+
for idx in squeeze_indices:
64+
data[idx] = np.squeeze(data[idx])
65+
return data
66+
67+
def inject(self, graph):
68+
for layer_name, data in self.params:
69+
if layer_name in graph:
70+
node = graph.get_node(layer_name)
71+
node.data = self.adjust_parameters(node, data)
72+
else:
73+
print_stderr('Ignoring parameters for non-existent layer: %s' % layer_name)
74+
75+
76+
class DataReshaper(object):
77+
78+
def __init__(self, mapping):
79+
self.mapping = mapping
80+
81+
def map(self, ndim):
82+
try:
83+
return self.mapping[ndim]
84+
except KeyError:
85+
raise KaffeError('Ordering not found for %d dimensional tensor.' % ndim)
86+
87+
def transpose(self, data):
88+
return data.transpose(self.map(data.ndim))
89+
90+
def has_spatial_parent(self, node):
91+
try:
92+
parent = node.get_only_parent()
93+
s = parent.output_shape
94+
return s.height > 1 or s.width > 1
95+
except KaffeError:
96+
return False
97+
98+
def reshape(self, graph, replace=True):
99+
for node in graph.nodes:
100+
if node.data is None:
101+
continue
102+
# Get the weights
103+
data = node.data[0]
104+
if (node.kind == NodeKind.InnerProduct) and self.has_spatial_parent(node):
105+
# The FC layer connected to the spatial layer needs to be
106+
# re-wired to match the new spatial ordering.
107+
in_shape = node.get_only_parent().output_shape
108+
fc_shape = ParamShape(data.shape)
109+
fc_order = self.map(2)
110+
data = data.reshape((fc_shape.output_channels, in_shape.channels, in_shape.height,
111+
in_shape.width))
112+
data = self.transpose(data)
113+
node.reshaped_data = data.reshape(fc_shape[fc_order[0]], fc_shape[fc_order[1]])
114+
else:
115+
node.reshaped_data = self.transpose(data)
116+
117+
if replace:
118+
for node in graph.nodes:
119+
if node.data is not None:
120+
# Set the weights
121+
node.data[0] = node.reshaped_data
122+
del node.reshaped_data

‎kaffe/errors.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import sys
2+
3+
class KaffeError(Exception):
4+
pass
5+
6+
def print_stderr(msg):
7+
sys.stderr.write('%s\n' % msg)

‎kaffe/core.py renamed to ‎kaffe/graph.py

Lines changed: 11 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
import numpy as np
21
from google.protobuf import text_format
32

4-
from .layers import *
5-
from .base import print_stderr
6-
from .caffe import has_pycaffe, get_caffe_resolver
3+
from .caffe import get_caffe_resolver
4+
from .data import DataInjector
5+
from .errors import KaffeError, print_stderr
6+
from .layers import LayerAdapter, LayerType, NodeKind, NodeDispatch
7+
from .shapes import TensorShape
78

89
class Node(object):
910

@@ -41,12 +42,6 @@ def parameters(self):
4142
return self.layer.parameters
4243
return None
4344

44-
@property
45-
def data_shape(self):
46-
assert self.data
47-
# pylint: disable=unsubscriptable-object
48-
return self.data[IDX_WEIGHTS].shape
49-
5045
def __str__(self):
5146
return '[%s] %s' % (self.kind, self.name)
5247

@@ -102,138 +97,25 @@ def visit(node):
10297
def compute_output_shapes(self):
10398
sorted_nodes = self.topologically_sorted()
10499
for node in sorted_nodes:
105-
node.output_shape = NodeKind.compute_output_shape(node)
100+
node.output_shape = TensorShape(*NodeKind.compute_output_shape(node))
106101

107102
def __contains__(self, key):
108103
return key in self.node_lut
109104

110105
def __str__(self):
111106
hdr = '{:<20} {:<30} {:>20} {:>20}'.format('Type', 'Name', 'Param', 'Output')
112107
s = [hdr, '-' * 94]
108+
fmt_tensor_shape = lambda ts: '({}, {}, {}, {})'.format(*ts)
113109
for node in self.topologically_sorted():
114-
data_shape = node.data[IDX_WEIGHTS].shape if node.data else '--'
110+
# If the node has learned parameters, display the first one's shape.
111+
# In case of convolutions, this corresponds to the weights.
112+
data_shape = node.data[0].shape if node.data else '--'
115113
out_shape = node.output_shape or '--'
116114
s.append('{:<20} {:<30} {:>20} {:>20}'.format(node.kind, node.name, data_shape,
117-
out_shape))
115+
fmt_tensor_shape(out_shape)))
118116
return '\n'.join(s)
119117

120118

121-
class DataInjector(object):
122-
123-
def __init__(self, def_path, data_path):
124-
self.def_path = def_path
125-
self.data_path = data_path
126-
self.did_use_pb = False
127-
self.params = None
128-
self.load()
129-
130-
def load(self):
131-
if has_pycaffe():
132-
self.load_using_caffe()
133-
else:
134-
self.load_using_pb()
135-
136-
def load_using_caffe(self):
137-
caffe = get_caffe_resolver().caffe
138-
net = caffe.Net(self.def_path, self.data_path, caffe.TEST)
139-
data = lambda blob: blob.data
140-
self.params = [(k, map(data, v)) for k, v in net.params.items()]
141-
142-
def load_using_pb(self):
143-
data = get_caffe_resolver().NetParameter()
144-
data.MergeFromString(open(self.data_path, 'rb').read())
145-
pair = lambda layer: (layer.name, self.transform_data(layer))
146-
layers = data.layers or data.layer
147-
self.params = [pair(layer) for layer in layers if layer.blobs]
148-
self.did_use_pb = True
149-
150-
def transform_data(self, layer):
151-
transformed = []
152-
for blob in layer.blobs:
153-
if len(blob.shape.dim):
154-
dims = blob.shape.dim
155-
c_o, c_i, h, w = map(int, [1] * (4 - len(dims)) + list(dims))
156-
else:
157-
c_o = blob.num
158-
c_i = blob.channels
159-
h = blob.height
160-
w = blob.width
161-
data = np.array(blob.data, dtype=np.float32).reshape(c_o, c_i, h, w)
162-
transformed.append(data)
163-
return transformed
164-
165-
def adjust_parameters(self, node, data):
166-
if not self.did_use_pb:
167-
return data
168-
# When using the protobuf-backend, each parameter initially has four dimensions.
169-
# In certain cases (like FC layers), we want to eliminate the singleton dimensions.
170-
# This implementation takes care of the common cases. However, it does leave the
171-
# potential for future issues.
172-
# The Caffe-backend does not suffer from this problem.
173-
data = list(data)
174-
squeeze_indices = [1] # Squeeze biases.
175-
if node.kind == NodeKind.InnerProduct:
176-
squeeze_indices.append(0) # Squeeze FC.
177-
for idx in squeeze_indices:
178-
data[idx] = np.squeeze(data[idx])
179-
return data
180-
181-
def inject(self, graph):
182-
for layer_name, data in self.params:
183-
if layer_name in graph:
184-
node = graph.get_node(layer_name)
185-
node.data = self.adjust_parameters(node, data)
186-
else:
187-
print_stderr('Ignoring parameters for non-existent layer: %s' % layer_name)
188-
189-
190-
class DataReshaper(object):
191-
192-
def __init__(self, mapping):
193-
self.mapping = mapping
194-
195-
def map(self, ndim):
196-
try:
197-
return self.mapping[ndim]
198-
except KeyError:
199-
raise KaffeError('Ordering not found for %d dimensional tensor.' % ndim)
200-
201-
def transpose(self, data):
202-
return data.transpose(self.map(data.ndim))
203-
204-
def has_spatial_parent(self, node):
205-
try:
206-
parent = node.get_only_parent()
207-
s = parent.output_shape
208-
return s[IDX_H] > 1 or s[IDX_W] > 1
209-
except KaffeError:
210-
return False
211-
212-
def reshape(self, graph, replace=True):
213-
for node in graph.nodes:
214-
if node.data is None:
215-
continue
216-
data = node.data[IDX_WEIGHTS]
217-
if (node.kind == NodeKind.InnerProduct) and self.has_spatial_parent(node):
218-
# The FC layer connected to the spatial layer needs to be
219-
# re-wired to match the new spatial ordering.
220-
in_shape = node.get_only_parent().output_shape
221-
fc_shape = data.shape
222-
fc_order = self.map(2)
223-
data = data.reshape((fc_shape[IDX_C_OUT], in_shape[IDX_C], in_shape[IDX_H],
224-
in_shape[IDX_W]))
225-
data = self.transpose(data)
226-
node.reshaped_data = data.reshape(fc_shape[fc_order[0]], fc_shape[fc_order[1]])
227-
else:
228-
node.reshaped_data = self.transpose(data)
229-
230-
if replace:
231-
for node in graph.nodes:
232-
if node.data is not None:
233-
node.data[IDX_WEIGHTS] = node.reshaped_data
234-
del node.reshaped_data
235-
236-
237119
class GraphBuilder(object):
238120
'''Constructs a model graph from a Caffe protocol buffer definition.'''
239121

‎kaffe/layers.py

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -51,14 +51,7 @@
5151

5252
LAYER_TYPES = LAYER_DESCRIPTORS.keys()
5353

54-
55-
def generate_layer_type_enum():
56-
types = {t: t for t in LAYER_TYPES}
57-
return type('LayerType', (), types)
58-
59-
60-
LayerType = generate_layer_type_enum()
61-
54+
LayerType = type('LayerType', (), {t: t for t in LAYER_TYPES})
6255

6356
class NodeKind(LayerType):
6457

0 commit comments

Comments
 (0)