def get_all_linear_layers(model):
layers = []
all_layers = get_all_layers(model)
for idx, layer in enumerate(all_layers):
# check if linear layer
if hasattr(layer, 'kernel') or hasattr(layer, 'depthwise_kernel'):
if layer.activation.__name__ != "linear":
layers.append(layer)
else:
next_layer = all_layers[
idx + 1] if idx + 1 < len(all_layers) else None
if isinstance(next_layer, BatchNormalization):
next_layer = all_layers[
idx + 2] if idx + 2 < len(all_layers) else None
if isinstance(next_layer, Activation) or isinstance(
next_layer, ActivationQ):
layers.append(layer)
return layers
def transform(model, bits_w, bits_x, log=False, quantize=True, verif_preproc=False,
slalom=False, slalom_integrity=False, slalom_privacy=False, sgxutils=None, queues=None):
if slalom:
assert(quantize)
old_ops = K.get_session().graph.get_operations()
#all_layers = [[l] if not isinstance(l, ResNetBlock) else l.get_layers() for l in model.layers]
#all_layers = list(itertools.chain.from_iterable(all_layers))
all_layers = get_all_layers(model)
fuse_bn(all_layers)
queue_ctr = 0
layers = model.layers
layer_map = {}
flattened = False
def transform_layer(layer, next_layer, queue_ctr, flattened):
print("transform {} (next = {})".format(layer, next_layer))
new_layers = []
skip_next = False
if isinstance(layer, InputLayer):
new_layers.append(InputLayer.from_config(layer.get_config()))
elif isinstance(layer, Conv2D) and not isinstance(layer, DepthwiseConv2D):
conf = layer.get_config()
act = conf['activation']
# if the next layer is a pooling layer, create a fused activation
maxpool_params = None
if slalom and isinstance(next_layer, MaxPooling2D):
mp = next_layer
assert (layer.activation == relu)
maxpool_params = mp.get_config()
skip_next = True
act_layer = None
if act != "linear":
conf['activation'] = "linear"
if slalom and isinstance(next_layer, GlobalAveragePooling2D):
assert layer.activation in [relu, relu6]
act = "avgpool" + act
skip_next = True
act_layer = ActivationQ(act, bits_w, bits_x, maxpool_params=maxpool_params, log=log,
quantize=quantize, slalom=slalom, slalom_integrity=slalom_integrity,
slalom_privacy=slalom_privacy, sgxutils=sgxutils,
queue=None if queues is None else queues[queue_ctr])
queue_ctr += 1
conf['bits_w'] = bits_w
conf['bits_x'] = bits_x
conf['log'] = log
conf['quantize'] = quantize
conf['slalom'] = slalom
conf['slalom_integrity'] = slalom_integrity
conf['slalom_privacy'] = slalom_privacy
conf['sgxutils'] = sgxutils
new_layer = Conv2DQ.from_config(conf)
new_layers.append(new_layer)
layer_map[new_layer] = layer
if act_layer is not None:
new_layers.append(act_layer)
elif isinstance(layer, DepthwiseConv2D):
conf = layer.get_config()
assert conf['activation'] == "linear"
conf['bits_w'] = bits_w
conf['bits_x'] = bits_x
conf['log'] = log
conf['quantize'] = quantize
conf['slalom'] = slalom
conf['slalom_integrity'] = slalom_integrity
conf['slalom_privacy'] = slalom_privacy
conf['sgxutils'] = sgxutils
new_layer = DepthwiseConv2DQ.from_config(conf)
new_layers.append(new_layer)
layer_map[new_layer] = layer
elif isinstance(layer, Dense):
conf = layer.get_config()
act = conf['activation']
act_layer = None
if act != "linear":
conf['activation'] = "linear"
act_layer = ActivationQ(act, bits_w, bits_x, log=log,
quantize=quantize, slalom=slalom, slalom_integrity=slalom_integrity,
slalom_privacy=slalom_privacy, sgxutils=sgxutils,
queue=None if queues is None else queues[queue_ctr])
queue_ctr += 1
conf['bits_w'] = bits_w
conf['bits_x'] = bits_x
conf['log'] = log
conf['quantize'] = quantize
conf['slalom'] = slalom
conf['slalom_integrity'] = slalom_integrity
conf['slalom_privacy'] = slalom_privacy
conf['sgxutils'] = sgxutils
# replace the dense layer by a pointwise convolution
if verif_preproc:
del conf['units']
conf['filters'] = layer.units
conf['kernel_size'] = 1
if not flattened:
h_in = int(layer.input_spec.axes[-1])
new_layers.append(Reshape((1, 1, h_in)))
flattened = True
new_layer = Conv2DQ.from_config(conf)
new_layers.append(new_layer)
layer_map[new_layer] = layer
else:
new_layer = DenseQ.from_config(conf)
new_layers.append(new_layer)
layer_map[new_layer] = layer
if act_layer is not None:
new_layers.append(act_layer)
elif isinstance(layer, BatchNormalization):
pass
elif isinstance(layer, MaxPooling2D):
assert (not slalom or not slalom_privacy)
new_layers.append(MaxPooling2D.from_config(layer.get_config()))
elif isinstance(layer, AveragePooling2D):
assert (not slalom or not slalom_privacy)
new_layers.append(AveragePooling2D.from_config(layer.get_config()))
new_layers.append(Lambda(lambda x: K.round(x)))
elif isinstance(layer, Activation):
assert layer.activation in [relu6, relu, softmax]
queue = None if queues is None else queues[queue_ctr]
queue_ctr += 1
act_func = "relu6" if layer.activation == relu6 else "relu" if layer.activation == relu else "softmax"
if slalom and isinstance(next_layer, GlobalAveragePooling2D):
#assert layer.activation == relu6
act_func = "avgpoolrelu6"
skip_next = True
maxpool_params = None
if slalom and (isinstance(next_layer, MaxPooling2D) or isinstance(next_layer, AveragePooling2D)):
mp = next_layer
assert (layer.activation == relu)
maxpool_params = mp.get_config()
skip_next = True
new_layers.append(ActivationQ(act_func, bits_w, bits_x, log=log,
maxpool_params=maxpool_params,
quantize=quantize, slalom=slalom,
slalom_integrity=slalom_integrity,
slalom_privacy=slalom_privacy,
sgxutils=sgxutils, queue=queue))
elif isinstance(layer, ZeroPadding2D):
if quantize:
# merge with next layer
conv = next_layer
assert isinstance(conv, Conv2D) or isinstance(conv, DepthwiseConv2D)
assert conv.padding == 'valid'
conv.padding = 'same'
else:
new_layers.append(ZeroPadding2D.from_config(layer.get_config()))
elif isinstance(layer, Flatten):
if not verif_preproc:
new_layers.append(Flatten.from_config(layer.get_config()))
elif isinstance(layer, GlobalAveragePooling2D):
assert not slalom
conf = layer.get_config()
conf['bits_w'] = bits_w
conf['bits_x'] = bits_x
conf['log'] = log
conf['quantize'] = quantize
new_layers.append(GlobalAveragePooling2DQ.from_config(conf))
elif isinstance(layer, Reshape):
new_layers.append(Reshape.from_config(layer.get_config()))
elif isinstance(layer, Dropout):
pass
elif isinstance(layer, ResNetBlock):
#assert not slalom
path1 = []
path2 = []
for l in layer.path1:
lq, queue_ctr, _, _ = transform_layer(l, None, queue_ctr, flattened)
path1.extend(lq)
for l in layer.path2:
lq, queue_ctr, _, _ = transform_layer(l, None, queue_ctr, flattened)
path2.extend(lq)
[actq], queue_ctr, flattened, skip_next = transform_layer(layer.merge_act, next_layer, queue_ctr, flattened)
new_layer = ResNetBlock(layer.kernel_size, layer.filters, layer.stage, layer.block, layer.identity,
layer.strides, path1=path1, path2=path2, merge_act=actq,
quantize=quantize, bits_w=bits_w, bits_x=bits_x,
slalom=slalom, slalom_integrity=slalom_integrity, slalom_privacy=slalom_privacy)
new_layers.append(new_layer)
else:
raise AttributeError("Don't know how to handle layer {}".format(layer))
return new_layers, queue_ctr, flattened, skip_next
new_model = Sequential()
skip_next = False
l_count = 0
localf = False
while layers:
layer = layers.pop(0)
next_layer = layers[0] if len(layers) else None
if (l_count <= 6):
if not skip_next:
new_layers, queue_ctr, flattened, skip_next = transform_layer(layer, next_layer, queue_ctr, flattened)
for new_layer in new_layers:
new_model.add(new_layer)
else:
skip_next = False
else:
skip_next = False
if isinstance(layer, MaxPooling2D):
conf = layer.get_config()
new_layer = MaxPooling2D.from_config(conf)
new_layers.append(new_layer)
elif isinstance(layer, Flatten):
conf = layer.get_config()
new_layer = Flatten.from_config(conf)
new_layers.append(new_layer)
elif isinstance(layer, Conv2D):
conf = layer.get_config()
#conf['bits_w'] = bits_w
#conf['bits_x'] = bits_x
#conf["log"] = log
#conf['quantize'] = quantize
#conf['slalom'] = None
#conf['slalom_integrity'] = None
#conf['slalom_privacy'] = None
#conf['sgxutils'] = None
new_layer = layer
new_layers.append(new_layer)
elif isinstance(layer, Dense):
conf = layer.get_config()
#conf["log"] = log
#conf['quantize'] = quantize
#conf['slalom'] = None
#conf['slalom_integrity'] = None
#conf['slalom_privacy'] = None
#conf['sgxutils'] = None
#conf['bits_w'] = bits_w
#conf['bits_x'] = bits_x
new_layer = layer
new_layers.append(new_layer)
else:
print(type(layer))
assert(False)
new_model.add(new_layer)
l_count = l_count + 1
print("transformed summery")
print(new_model.summary())
# copy over (and potentially quantize) all the weights
new_layers = get_all_layers(new_model)
for layer in new_layers:
if layer in layer_map:
src_layer = layer_map[layer]
weights = src_layer.get_weights()
kernel = weights[0]
bias = weights[1]
print(layer)
if quantize and isinstance(layer, Conv2DQ):
range_w = 2**bits_w
range_x = 2**bits_x
kernel_q = np.round(range_w * kernel)
bias_q = np.round(range_w * range_x * bias)
if slalom_privacy:
if isinstance(layer, DepthwiseConv2DQ):
bias_q = bias_q.astype(np.float64)
kernel_q = kernel_q.astype(np.float64)
layer._trainable_weights = layer._trainable_weights[2:]
if isinstance(src_layer, Dense) and verif_preproc:
kernel_q = np.reshape(kernel_q, (1, 1, kernel_q.shape[0], kernel_q.shape[1]))
layer.set_weights((kernel_q, bias_q))
else:
layer._trainable_weights = layer._trainable_weights[2:]
layer.set_weights((kernel, bias))
# find all the TensorFlow ops that correspond to inputs/outputs of linear operators
new_ops = [op for op in K.get_session().graph.get_operations() if op not in old_ops]
linear_ops_in = [tf.reshape(op.inputs[0], [-1]) for op in new_ops if op.type in ['Conv2D', 'MatMul', 'DepthwiseConv2dNative']]
linear_ops_out = [tf.reshape(op.outputs[0], [-1]) for op in new_ops if op.type in ['BiasAdd']]
return new_model, linear_ops_in, linear_ops_out