def to_bitmap(shape, attr, contrib):
mask = np.zeros(np.prod(shape), dtype=np.int8)
pos_attr = attr[contrib > 0]
mask[TraceKey.to_array(pos_attr)] = 1
neg_attr = attr[contrib < 0]
mask[TraceKey.to_array(neg_attr)] = -1
return mask.reshape(shape)
def reconstruct_edge_from_trace(
trace,
graph,
node_name,
key=TraceKey.EDGE,
):
attrs = trace.nodes[node_name]
op = graph.op(graph.id(node_name))
if key not in attrs:
return None
else:
attr = attrs[key]
edge = TraceKey.to_array(attr)
return edge
def get_edge_mask(op: Union[Conv2dOp, PoolOp], compact: bool) -> np.ndarray:
edge_shape = op.attrs[TraceKey.EDGE_SHAPE]
if not compact:
return TraceKey.to_mask(op.attrs[TraceKey.EDGE],
(np.prod(edge_shape[:3], edge_shape[3:])),
compact)
input_tensor: Tensor = op.input_nodes[0]
output_tensor: Tensor = op.output_nodes[0]
edge = TraceKey.to_array(op.attrs[TraceKey.EDGE], compact)
input_shape = input_tensor.attrs[TraceKey.POINT_SHAPE]
output_shape = output_tensor.attrs[TraceKey.POINT_SHAPE]
if op.data_format == "NHWC":
input_shape = (input_shape[2], input_shape[0], input_shape[1])
output_shape = (output_shape[2], output_shape[0], output_shape[1])
if isinstance(op, Conv2dOp):
in_channel, kernel_height, kernel_width, out_channel, out_height, out_width = np.unravel_index(
edge, edge_shape)
else:
kernel_height, kernel_width, out_channel, out_height, out_width = np.unravel_index(
edge, edge_shape)
stride = np.array(op.strides)
kernel_size = (np.array(op.attrs[TraceKey.WEIGHT_SHAPE])[2:] if isinstance(
op, Conv2dOp) else np.array(op.filter_shape))
padding = calc_padding(
np.array(input_shape)[1:],
np.array(output_shape)[1:], stride, kernel_size)
in_height = kernel_height + (out_height * stride[0]) - padding[1][0]
in_width = kernel_width + (out_width * stride[1]) - padding[2][0]
edge_output_index = np.ravel_multi_index(
(out_channel, out_height, out_width), output_shape)
if isinstance(op, Conv2dOp):
edge_input_index = np.ravel_multi_index(
(in_channel, in_height, in_width), input_shape)
else:
edge_input_index = np.ravel_multi_index(
(out_channel, in_height, in_width), input_shape)
mask = np.zeros((np.prod(input_shape), np.prod(output_shape)),
dtype=np.int8)
mask[(edge_input_index, edge_output_index)] = 1
return mask
def after_create_session(self, session, coord):
tf_graph = tf.get_default_graph()
graph = self.graph
variables = get_variables_from_tf(graph, tf_graph, session)
for op in graph.ops:
if TraceKey.WEIGHT in op.attrs:
weight_name = get_weight(op).name
weight = variables[weight_name]
weight_shape_in_trace = op.attrs[TraceKey.WEIGHT_SHAPE]
traced_weight = np.unravel_index(
TraceKey.to_array(op.attrs[TraceKey.WEIGHT]),
weight_shape_in_trace)
if isinstance(op, Conv2dOp):
traced_weight = tuple(
[traced_weight[axis] for axis in [2, 3, 1, 0]])
elif isinstance(op, DenseOp):
traced_weight = tuple(
[traced_weight[axis] for axis in [1, 0]])
else:
raise RuntimeError()
mask = np.ones(weight.shape, dtype=np.int32)
mask[traced_weight] = 0
weight[mask.astype(np.bool)] = 0
load_variables_into_tf(tf_graph, variables, session)
def to_bitmap(shape, attr):
mask = np.zeros(np.prod(shape), dtype=np.int8)
mask[TraceKey.to_array(attr)] = 1
return mask.reshape(shape)
def reconstruct_edge(
trace,
graph,
key,
node_name,
):
attrs = trace.nodes[node_name]
op = graph.op(graph.id(node_name))
if key not in attrs:
return None
else:
attr = attrs[key]
edge = TraceKey.to_array(attr)
if isinstance(op, (AddOp, DenseOp)):
shape = attrs[key + "_shape"]
mask = np.zeros(np.prod(shape), dtype=np.int8)
mask[edge] = 1
mask = np.reshape(mask, shape)
return mask
elif isinstance(op, (MaxPoolOp, Conv2dOp, AvgPoolOp)):
input_shape = trace.tensors[op.input_nodes[0].name][
TraceKey.POINT_SHAPE]
output_shape = trace.tensors[op.output_nodes[0].name][
TraceKey.POINT_SHAPE]
if op.data_format == "NHWC":
input_shape = (
input_shape[2],
input_shape[0],
input_shape[1],
)
output_shape = (
output_shape[2],
output_shape[0],
output_shape[1],
)
in_channel, in_height, in_width, out_channel, out_height, out_width = np.unravel_index(
edge, input_shape + output_shape)
stride = np.array(op.strides)
kernel_size = (np.array(attrs[TraceKey.WEIGHT_SHAPE])[2:]
if isinstance(op, Conv2dOp) else np.array(
op.filter_shape))
padding = calc_padding(
np.array(input_shape)[1:],
np.array(output_shape)[1:],
stride,
kernel_size,
)
kernel_height = (in_height + padding[1][0] -
(out_height * stride[0]))
kernel_width = (in_width + padding[2][0] - (out_width * stride[1]))
edge_shape = attrs[TraceKey.EDGE_SHAPE]
if isinstance(op, Conv2dOp):
new_edge_index = np.ravel_multi_index(
(
in_channel,
kernel_height,
kernel_width,
out_channel,
out_height,
out_width,
),
edge_shape,
)
else:
new_edge_index = np.ravel_multi_index(
(
kernel_height,
kernel_width,
out_channel,
out_height,
out_width,
),
edge_shape,
)
mask = np.zeros(np.prod(edge_shape), dtype=np.int8)
mask[new_edge_index] = 1
mask = np.reshape(mask, edge_shape)
return mask
def to_bitmap(shape):
mask = np.zeros(np.prod(shape), dtype=np.int8)
mask[TraceKey.to_array(attr)] = 1
return np.packbits(mask)