def __call__(self, graph_name, value):
if isinstance(value.data, torch.Tensor):
nndct_tensor = Tensor(
name=get_full_name(graph_name, value.name),
shape=value.shape,
dtype=value.dtype,
data=value.data.cpu().numpy())
else:
nndct_tensor = Tensor(
name=get_full_name(graph_name, value.name), shape=value.shape, dtype=value.dtype)
return nndct_tensor
def DecoupleSharedParamsInConv(self):
# decouple shared parameters in graph
bias_tensor_list = []
weight_tensor_list = []
for node in self._graph.nodes:
if node.op.type in _CONV_TYPES:
weight_tensor = node.op.params[node.op.ParamName.WEIGHTS]
weight_name = weight_tensor.name
if weight_name in weight_tensor_list:
node_idx = node.idx
weight_name_copy = weight_name + '.' + str(node_idx)
new_weight_tensor = Tensor(name=weight_name_copy)
new_weight_tensor.clone_from(weight_tensor)
node.op.set_param(node.op.ParamName.WEIGHTS,
new_weight_tensor)
else:
weight_tensor_list.append(weight_name)
if node.node_attr(node.op.AttrName.BIAS_TERM):
bias_tensor = node.op.params[node.op.ParamName.BIAS]
bias_name = bias_tensor.name
if bias_name in bias_tensor_list:
node_idx = node.idx
bias_name_copy = bias_name + '.' + str(node_idx)
new_bias_tensor = Tensor(name=bias_name_copy)
new_bias_tensor.clone_from(bias_tensor)
node.op.set_param(node.op.ParamName.BIAS,
new_bias_tensor)
else:
bias_tensor_list.append(bias_name)
def __call__(self, scope, value):
if isinstance(value.data, torch.Tensor):
nndct_tensor = Tensor(name=get_full_name(scope, value.name),
shape=value.shape,
dtype=value.dtype,
data=value.data.cpu().numpy(),
layout=value.layout)
else:
nndct_tensor = Tensor(name=get_full_name(scope, value.name),
shape=value.shape,
dtype=value.dtype,
data=value.data,
layout=value.layout)
return nndct_tensor
def _generate_scale_node(node, scale_param, node_idx):
node_name = node.name
tensor_out = node.out_tensors[0]
out_tensor_shape = tensor_out.shape
scale_op = TorchBaseOperation(NNDCT_OP.CHANNEL_SCALE,
"Channel_Scale",
force_to_primitive=True)
scale_op.set_config("channel_scale", scale_param)
scale_op.set_config("input", tensor_out)
if node.op.type == NNDCT_OP.CHANNEL_SCALE:
scale_node_name = node_name[:node_name.rfind(".")] + ".2"
else:
scale_node_name = "/".join([node_name, "channel_scale.1"])
scale_node = Node(scale_node_name,
op=scale_op,
dtype="float32",
idx=node_idx,
in_quant_part=False)
out_tensor = Tensor(name=f"{scale_node_name}.0",
node=scale_node,
shape=out_tensor_shape,
dtype="float32",
layout=tensor_out.layout)
scale_node.out_tensors.append(out_tensor)
scale_node.in_tensors.append(tensor_out)
return scale_node
def __call__(self,
graph,
node_name,
op,
num_out_tensors,
shape=None,
in_tensors=None,
in_quant_part=True):
node_name = get_full_name(graph.name, node_name)
node = Node(node_name,
op=op,
dtype="float32",
idx=self._idx,
in_quant_part=in_quant_part)
# print(f"{node.name} quant state: {node.in_quant_part}")
for i in range(num_out_tensors):
tensor = Tensor(name=f"{node_name}_{i}", node=node, shape=shape)
node.out_tensors.append(tensor)
if in_tensors:
for tensor in in_tensors:
node.in_tensors.append(tensor)
graph.add_node(node)
self._idx += 1
def _convert_tensor(self, value, scope=None):
if scope is None:
assert self.cur_graph
value_scope = self.cur_graph.name
else:
value_scope = scope
if isinstance(value.data, torch.Tensor):
nndct_tensor = Tensor(name=get_full_name(value_scope, value.name),
shape=value.shape,
dtype=value.dtype,
data=value.data.cpu().numpy(),
layout=value.layout)
else:
nndct_tensor = Tensor(name=get_full_name(value_scope, value.name),
shape=value.shape,
dtype=value.dtype,
data=value.data,
layout=value.layout)
return nndct_tensor
def _load_data(self, module):
for node in self._nndct_graph.nodes:
if node.op.type in [NNDCT_OP.BASIC_LSTM, NNDCT_OP.BASIC_GRU]:
for nndct_param, param_tensors in node.op.params.items():
for tensor in param_tensors:
data = module.state_dict()[get_short_name(
tensor.name)].cpu().numpy()
tensor.from_ndarray(data)
tensor = tensor_util.convert_parameter_tensor_format(
tensor, FrameworkType.TORCH, FrameworkType.NNDCT)
#combine bias_ih and bias_hh item
if node.op.type == NNDCT_OP.BASIC_LSTM:
for bias_term in [
node.op.ParamName.BIAS,
node.op.ParamName.BIAS_REVERSE
]:
if bias_term in node.op.params and len(
node.op.params[bias_term]) > 0:
if len(node.op.params[bias_term]) % 2 != 0:
raise RuntimeError(
"The num of bias should be even")
i = 0
bias_list = []
while i != len(node.op.params[bias_term]):
bias_ih = node.op.params[bias_term][i]
bias_hh = node.op.params[bias_term][i + 1]
tensor_name = f"bias_{i//2}" if bias_term == node.op.ParamName.BIAS else f"bias_{i//2}_reverse"
bias = Tensor(name=get_full_name(
self._nndct_graph.name, tensor_name),
data=bias_ih.data + bias_hh.data)
bias_list.append(bias)
i = i + 2
node.op.set_param(bias_term, bias_list)
elif node.op.type == NNDCT_OP.CONVTRANSPOSE2D:
for param_name, tensor in node.op.params.items():
data = module.state_dict()[get_short_name(
tensor.name)].cpu().numpy()
if param_name == node.op.ParamName.WEIGHTS:
data = np.copy(data).transpose(1, 0, 2, 3)
data = np.ascontiguousarray(data)
tensor.from_ndarray(data)
tensor = tensor_util.convert_parameter_tensor_format(
tensor, FrameworkType.TORCH, FrameworkType.NNDCT)
elif node.op.type == NNDCT_OP.DEPTHWISE_CONV2D:
for param_name, tensor in node.op.params.items():
data = module.state_dict()[get_short_name(
tensor.name)].cpu().numpy()
if param_name == node.op.ParamName.WEIGHTS:
in_channels = node.node_config("in_channels")
out_channels = node.node_config("out_channels")
kernel_size = node.node_config("kernel_size")
channel_mutiplier = int(out_channels / in_channels)
data = np.copy(data).reshape(
(channel_mutiplier, in_channels, *kernel_size))
tensor.from_ndarray(data)
tensor = tensor_util.convert_parameter_tensor_format(
tensor, FrameworkType.TORCH, FrameworkType.NNDCT)
else:
for param_name, tensor in node.op.params.items():
data = module.state_dict()[get_short_name(
tensor.name)].cpu().numpy()
tensor.from_ndarray(data)
tensor = tensor_util.convert_parameter_tensor_format(
tensor, FrameworkType.TORCH, FrameworkType.NNDCT)
def basic_lstm(self, node):
graph_scope_name = node.name.split(_GRAPH_SCOPE_SYM)[0]
node_creator = _NodeCreator()
graphs = []
bidirectional = node.node_attr(node.op.AttrName.BIDIRECTIONAL)
lstm_direction = ["forward"]
if bidirectional:
lstm_direction = ["forward", "backward"]
for i in range(node.node_attr(node.op.AttrName.NUM_LAYERS)):
lstm_cell_pair = {}
if i == 0:
input_size = node.node_attr(node.op.AttrName.INPUT_SIZE)
else:
input_size = len(lstm_direction) * node.node_attr(
node.op.AttrName.HIDDEN_SIZE)
hidden_size = node.node_attr(node.op.AttrName.HIDDEN_SIZE)
bias=True
for direction in lstm_direction:
if direction == "forward":
w_ih = node.op.params[node.op.ParamName.WEIGHT_IH][i]
w_hh = node.op.params[node.op.ParamName.WEIGHT_HH][i]
if node.op.ParamName.BIAS in node.op.params:
bias_hi = node.op.params[node.op.ParamName.BIAS][i]
else:
bias=False
else:
w_ih = node.op.params[node.op.ParamName.WEIGHT_IH_REVERSE][i]
w_hh = node.op.params[node.op.ParamName.WEIGHT_HH_REVERSE][i]
if node.op.ParamName.BIAS_REVERSE in node.op.params:
bias_hi = node.op.params[node.op.ParamName.BIAS_REVERSE][i]
else:
bias=False
# lstm_node_name = node.name.replace("/", "_")
graph_name = f"{graph_scope_name}_StandardLstmCell_layer_{i}_{direction}"
graph = Graph(graph_name=graph_name)
lstm_cell_pair[direction] = graph
w_ii = Tensor(get_full_name(graph.name, "weight_ii"))
w_if = Tensor(get_full_name(graph.name, "weight_if"))
w_ig = Tensor(get_full_name(graph.name, "weight_ig"))
w_io = Tensor(get_full_name(graph.name, "weight_io"))
w_ii.from_ndarray(w_ih.data[:hidden_size])
w_if.from_ndarray(w_ih.data[hidden_size:2 * hidden_size])
w_ig.from_ndarray(w_ih.data[2 * hidden_size:3 * hidden_size])
w_io.from_ndarray(w_ih.data[3 * hidden_size:4 * hidden_size])
w_hi = Tensor(get_full_name(graph.name, "weight_hi"))
w_hf = Tensor(get_full_name(graph.name, "weight_hf"))
w_hg = Tensor(get_full_name(graph.name, "weight_hg"))
w_ho = Tensor(get_full_name(graph.name, "weight_ho"))
w_hi.from_ndarray(w_hh.data[:hidden_size])
w_hf.from_ndarray(w_hh.data[hidden_size:2 * hidden_size])
w_hg.from_ndarray(w_hh.data[2 * hidden_size:3 * hidden_size])
w_ho.from_ndarray(w_hh.data[3 * hidden_size:4 * hidden_size])
bias_i = Tensor(get_full_name(graph.name, "bias_i"))
bias_f = Tensor(get_full_name(graph.name, "bias_f"))
bias_g = Tensor(get_full_name(graph.name, "bias_g"))
bias_o = Tensor(get_full_name(graph.name, "bias_o"))
if bias is True:
bias_i.from_ndarray(bias_hi.data[:hidden_size])
bias_f.from_ndarray(bias_hi.data[hidden_size:2 * hidden_size])
bias_g.from_ndarray(bias_hi.data[2 * hidden_size:3 * hidden_size])
bias_o.from_ndarray(bias_hi.data[3 * hidden_size:4 * hidden_size])
op = TorchBaseOperation(NNDCT_OP.INPUT, NNDCT_OP.INPUT)
op.set_config("input", "args[0]")
shape = [1, input_size]
node_creator(
graph=graph,
node_name="input_0",
op=op,
num_out_tensors=1,
shape=shape)
op = TorchBaseOperation(NNDCT_OP.INPUT, NNDCT_OP.INPUT)
op.set_config("input", "args[1]")
shape = [1, hidden_size]
node_creator(
graph=graph,
node_name="input_1",
op=op,
num_out_tensors=1,
shape=shape)
op = TorchBaseOperation(NNDCT_OP.INPUT, NNDCT_OP.INPUT)
op.set_config("input", "args[2]")
shape = [1, hidden_size]
node_creator(
graph=graph,
node_name="input_2",
op=op,
num_out_tensors=1,
shape=shape)
# y_i = w_ii * input_0 + bias_i + w_hi * input_1
op = TorchLinear()
op.set_config("bias", bias)
op.set_config("out_features", hidden_size)
op.set_config("in_features", input_size)
op.set_param(op.ParamName.WEIGHTS, w_ii)
if bias is True:
op.set_param(op.ParamName.BIAS, bias_i)
node_creator(
graph=graph,
node_name="w_ii * input_0 + bias_i",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_0")).out_tensors)
op = TorchLinear()
op.set_config("bias", False)
op.set_config("out_features", hidden_size)
op.set_config("in_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_hi)
# op.set_param(op.ParamName.BIAS, bias_i)
node_creator(
graph=graph,
node_name="w_hi * input_1",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_1")).out_tensors)
op = TorchAdd()
op.set_config("input", graph.node(get_full_name(graph.name, "w_ii * input_0 + bias_i")).out_tensors[0])
op.set_config("other",
graph.node(get_full_name(graph.name, "w_hi * input_1")).out_tensors[0])
node_creator(
graph=graph,
node_name="y_i",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "w_ii * input_0 + bias_i")).out_tensors[0],
graph.node(get_full_name(graph.name, "w_hi * input_1")).out_tensors[0]
])
# y_f = w_if * input_0 + bias_f + w_hf * input_1
op = TorchLinear()
op.set_config("bias", bias)
op.set_config("in_features", input_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_if)
if bias is True:
op.set_param(op.ParamName.BIAS, bias_f)
node_creator(
graph=graph,
node_name="w_if * input_0 + bias_f",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_0")).out_tensors)
op = TorchLinear()
op.set_config("bias", False)
op.set_config("in_features", hidden_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_hf)
# op.set_param(op.ParamName.BIAS, bias_f)
node_creator(
graph=graph,
node_name="w_hf * input_1",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_1")).out_tensors)
op = TorchAdd()
op.set_config("input", graph.node(get_full_name(graph.name, "w_if * input_0 + bias_f")).out_tensors[0])
op.set_config("other",
graph.node(get_full_name(graph.name, "w_hf * input_1")).out_tensors[0])
node_creator(
graph=graph,
node_name="y_f",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "w_if * input_0 + bias_f")).out_tensors[0],
graph.node(get_full_name(graph.name, "w_hf * input_1")).out_tensors[0]
])
# y_g = w_ig * input_0 + bias_g + w_hg * input_1
op = TorchLinear()
op.set_config("bias", bias)
op.set_config("in_features", input_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_ig)
if bias is True:
op.set_param(op.ParamName.BIAS, bias_g)
node_creator(
graph=graph,
node_name="w_ig * input_0 + bias_g",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_0")).out_tensors)
op = TorchLinear()
op.set_config("bias", False)
op.set_config("in_features", hidden_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_hg)
# op.set_param(op.ParamName.BIAS, bias_g)
node_creator(
graph=graph,
node_name="w_hg * input_1",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_1")).out_tensors)
op = TorchAdd()
op.set_config("input", graph.node(get_full_name(graph.name, "w_ig * input_0 + bias_g")).out_tensors[0])
op.set_config("other",
graph.node(get_full_name(graph.name, "w_hg * input_1")).out_tensors[0])
node_creator(
graph=graph,
node_name="y_g",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "w_ig * input_0 + bias_g")).out_tensors[0],
graph.node(get_full_name(graph.name, "w_hg * input_1")).out_tensors[0]
])
# y_o = w_io * input_0 + bias_o + w_ho * input_1
op = TorchLinear()
op.set_config("bias", bias)
op.set_config("in_features", input_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_io)
if bias is True:
op.set_param(op.ParamName.BIAS, bias_o)
node_creator(
graph=graph,
node_name="w_io * input_0 + bias_o",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_0")).out_tensors)
op = TorchLinear()
op.set_config("bias", False)
op.set_config("in_features", hidden_size)
op.set_config("out_features", hidden_size)
op.set_param(op.ParamName.WEIGHTS, w_ho)
# op.set_param(op.ParamName.BIAS, bias_o)
node_creator(
graph=graph,
node_name="w_ho * input_1",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "input_1")).out_tensors)
op = TorchAdd()
op.set_config("input", graph.node(get_full_name(graph.name, "w_io * input_0 + bias_o")).out_tensors[0])
op.set_config("other",
graph.node(get_full_name(graph.name, "w_ho * input_1")).out_tensors[0])
node_creator(
graph=graph,
node_name="y_o",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "w_io * input_0 + bias_o")).out_tensors[0],
graph.node(get_full_name(graph.name, "w_ho * input_1")).out_tensors[0]
])
# op = Split(optype=NNDCT_OP.SPLIT)
# op.set_attr(op.AttrName.INPUT, graph.node("combine_2_linearity").out_tensors[0])
# op.set_attr(op.AttrName.SPLIT_SIZE_OR_SECTIONS, hidden_size)
# op.set_attr(op.AttrName.AXIS, 1)
# node_creator(graph=graph,
# node_name="split_ifgo",
# op=op,
# num_out_tensors=4,
# in_tensors=graph.node("combine_2_linearity").out_tensors)
op = TorchSigmoid()
node_creator(
graph=graph,
node_name="it",
op=op,
num_out_tensors=1,
in_tensors=[graph.node(get_full_name(graph.name, "y_i")).out_tensors[0]])
op = TorchSigmoid()
node_creator(
graph=graph,
node_name="ft",
op=op,
num_out_tensors=1,
in_tensors=[graph.node(get_full_name(graph.name, "y_f")).out_tensors[0]])
op = TorchTanh()
node_creator(
graph=graph,
node_name="cct",
op=op,
num_out_tensors=1,
in_tensors=[graph.node(get_full_name(graph.name, "y_g")).out_tensors[0]])
op = TorchSigmoid()
node_creator(
graph=graph,
node_name="ot",
op=op,
num_out_tensors=1,
in_tensors=[graph.node(get_full_name(graph.name, "y_o")).out_tensors[0]])
op = TorchMul()
op.set_config("input", graph.node(get_full_name(graph.name, "it")).out_tensors[0])
op.set_config("other", graph.node(get_full_name(graph.name, "cct")).out_tensors[0])
node_creator(
graph=graph,
node_name="it*cct",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "it")).out_tensors[0],
graph.node(get_full_name(graph.name, "cct")).out_tensors[0]
])
op = TorchMul()
op.set_config("input", graph.node(get_full_name(graph.name, "ft")).out_tensors[0])
op.set_config("other", graph.node(get_full_name(graph.name, "input_2")).out_tensors[0])
node_creator(
graph=graph,
node_name="ft*input_2",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "ft")).out_tensors[0],
graph.node(get_full_name(graph.name, "input_2")).out_tensors[0]
])
op = TorchAdd()
op.set_config("input", graph.node(get_full_name(graph.name, "it*cct")).out_tensors[0])
op.set_config("other", graph.node(get_full_name(graph.name, "ft*input_2")).out_tensors[0])
node_creator(
graph=graph,
node_name="c_next",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "it*cct")).out_tensors[0],
graph.node(get_full_name(graph.name, "ft*input_2")).out_tensors[0]
])
op = TorchTanh()
node_creator(
graph=graph,
node_name="c_temp",
op=op,
num_out_tensors=1,
in_tensors=graph.node(get_full_name(graph.name, "c_next")).out_tensors)
op = TorchMul()
op.set_config("input", graph.node(get_full_name(graph.name, "ot")).out_tensors[0])
op.set_config("other", graph.node(get_full_name(graph.name, "c_temp")).out_tensors[0])
node_creator(
graph=graph,
node_name="h_next",
op=op,
num_out_tensors=1,
in_tensors=[
graph.node(get_full_name(graph.name, "ot")).out_tensors[0],
graph.node(get_full_name(graph.name, "c_temp")).out_tensors[0]
])
self._connect_nodes(graph)
graph.add_end_tensor(graph.node(get_full_name(graph.name, "h_next")).out_tensors[0])
graph.add_end_tensor(graph.node(get_full_name(graph.name, "c_next")).out_tensors[0])
graphs.append(lstm_cell_pair)
return graphs
