def load_kalid_nnet2_model(graph, file_descr, nnet_name): input_name = 'Input' graph.add_node(input_name, name=input_name, kind='op', op='Parameter', parameters=None, shape=None) prev_layer_id = input_name all_components = load_components(file_descr, graph) used_layers = set() for layer_id in all_components: prev_node = Node(graph, prev_layer_id) if prev_node.op == 'Parameter': parameters = Node(graph, layer_id).parameters input_dim = read_token_value(parameters, b'<InputDim>') prev_node['shape'] = np.array([1, input_dim], dtype=np.int64) prev_node.add_output_port(0) Node(graph, layer_id).add_input_port(0) graph.create_edge(prev_node, Node(graph, layer_id), 0, 0, create_edge_attrs(prev_layer_id, layer_id, prev_layer_id)) used_layers.add(prev_layer_id) prev_layer_id = layer_id log.debug('{} and {} were connected'.format(prev_layer_id, layer_id)) # Tensor names information corresponding to a node is stored on outgoing edges. # As output nodes do not have outgoing edges, fake outputs are required. In the following code # for each output Identity node is added, and tensor name for the output is kept # on (output, fake output) edge. After Result nodes adding transformation fake outputs # are deleted from graph. output_layers = graph.nodes - used_layers add_outputs_identity(graph, output_layers, lambda g, output, fake_output: g.create_edge( Node(g, output), Node(g, fake_output), 0, 0, create_edge_attrs(output, fake_output, output)))
def load_kalid_nnet1_model(graph, file_descr, name): prev_layer_id = 'Parameter' graph.add_node(prev_layer_id, name=prev_layer_id, kind='op', op='Parameter', parameters=None) used_layers = set() while True: component_type = find_next_component(file_descr) if component_type == end_of_nnet_tag.lower()[1:-1]: break layer_o = read_binary_integer32_token(file_descr) layer_i = read_binary_integer32_token(file_descr) if component_type == 'parallelcomponent': prev_layer_id = load_parallel_component(file_descr, graph, prev_layer_id) find_end_of_component(file_descr, component_type) continue start_index = file_descr.tell() end_tag, end_index = find_end_of_component(file_descr, component_type) end_index -= len(end_tag) layer_id = graph.unique_id(prefix=component_type) graph.add_node(layer_id, parameters=get_parameters(file_descr, start_index, end_index), op=component_type, kind='op', layer_i=layer_i, layer_o=layer_o) prev_node = Node(graph, prev_layer_id) if prev_node.op == 'Parameter': prev_node['shape'] = np.array([1, layer_i], dtype=np.int64) prev_node.add_output_port(0) Node(graph, layer_id).add_input_port(0) graph.create_edge( prev_node, Node(graph, layer_id), 0, 0, create_edge_attrs(prev_layer_id, layer_id, prev_layer_id)) used_layers.add(prev_layer_id) prev_layer_id = layer_id log.debug('{} (type is {}) was loaded'.format(prev_layer_id, component_type)) # Tensor names information corresponding to a node is stored on outgoing edges. # As output nodes do not have outgoing edges, fake outputs are required. In the following code # for each output Identity node is added, and tensor name for the output is kept # on (output, fake output) edge. After Result nodes adding transformation fake outputs # are deleted from graph. output_layers = graph.nodes - used_layers add_outputs_identity( graph, output_layers, lambda g, output, fake_output: g.create_edge( Node(g, output), Node(g, fake_output), 0, 0, create_edge_attrs(output, fake_output, output)))
def load_kalid_nnet2_model(graph, file_descr, nnet_name): input_name = 'Input' graph.add_node(input_name, name=input_name, kind='op', op='Parameter', parameters=None, shape=None) prev_layer_id = input_name all_components = load_components(file_descr, graph) for layer_id in all_components: prev_node = Node(graph, prev_layer_id) if prev_node.op == 'Parameter': parameters = Node(graph, layer_id).parameters input_dim = read_token_value(parameters, b'<InputDim>') prev_node['shape'] = np.array([1, input_dim], dtype=np.int64) prev_node.add_output_port(0) Node(graph, layer_id).add_input_port(0) graph.create_edge( prev_node, Node(graph, layer_id), 0, 0, create_edge_attrs(prev_layer_id, layer_id, prev_layer_id)) prev_layer_id = layer_id log.debug('{} and {} were connected'.format(prev_layer_id, layer_id))
def load_kalid_nnet1_model(file_descr, name): graph = Graph(name=name) prev_layer_id = 'Input' graph.add_node(prev_layer_id, name=prev_layer_id, kind='op', op='Input', parameters=None) input_shape = [] while True: component_type = find_next_component(file_descr) if component_type == end_of_nnet_tag.lower()[1:-1]: break layer_o = read_binary_integer32_token(file_descr) layer_i = read_binary_integer32_token(file_descr) if component_type == 'parallelcomponent': prev_layer_id = load_parallel_component(file_descr, graph, prev_layer_id) continue start_index = file_descr.tell() end_tag, end_index = find_end_of_component(file_descr, component_type) end_index -= len(end_tag) layer_id = graph.unique_id(prefix=component_type) graph.add_node(layer_id, parameters=get_parameters(file_descr, start_index, end_index), op=component_type, kind='op', layer_i=layer_i, layer_o=layer_o) prev_node = Node(graph, prev_layer_id) if prev_node.op == 'Input': prev_node['shape'] = np.array([1, layer_i], dtype=np.int64) input_shape = np.array([1, layer_i], dtype=np.int64) graph.add_edge(prev_layer_id, layer_id, **create_edge_attrs(prev_layer_id, layer_id)) prev_layer_id = layer_id log.debug('{} (type is {}) was loaded'.format(prev_layer_id, component_type)) return graph, input_shape
def load_kalid_nnet2_model(file_descr, nnet_name): graph = Graph(name=nnet_name) input_name = 'Input' input_shape = np.array([]) graph.add_node(input_name, name=input_name, kind='op', op='Input', parameters=None, shape=None) prev_layer_id = input_name collect_until_token(file_descr, b'<Nnet>') num_components = read_token_value(file_descr, b'<NumComponents>') log.debug('Network contains {} components'.format(num_components)) collect_until_token(file_descr, b'<Components>') for _ in range(num_components): component_type = find_next_component(file_descr) if component_type == end_of_nnet_tag.lower()[1:-1]: break start_index = file_descr.tell() end_tag, end_index = find_end_of_component(file_descr, component_type) layer_id = graph.unique_id(prefix=component_type) graph.add_node(layer_id, parameters=get_parameters(file_descr, start_index, end_index), op=component_type, kind='op') prev_node = Node(graph, prev_layer_id) if prev_node.op == 'Input': parameters = Node(graph, layer_id).parameters input_dim = read_token_value(parameters, b'<InputDim>') prev_node['shape'] = np.array([1, input_dim], dtype=np.int64) input_shape = np.array([1, input_dim], dtype=np.int64) graph.add_edge(prev_layer_id, layer_id, **create_edge_attrs(prev_layer_id, layer_id)) prev_layer_id = layer_id log.debug('{} (type is {}) was loaded'.format(prev_layer_id, component_type)) return graph, input_shape
def parse_specifier(string, graph, layer_node_map): pos = string.find(b'(') if pos == -1: # node name input_name = str(string.split(b' ')[0]).strip('b').replace( "\'", '').replace('\\n', '') if input_name not in layer_node_map: node_name = graph.unique_id(prefix=input_name) graph.add_node(node_name, parameters=[], op="", kind='op') layer_node_map[input_name] = node_name else: node_name = layer_node_map[input_name] return node_name spec = string[:pos] args = get_args_for_specifier(string[pos:]) if spec == b'Append': nodes = [] for i in range(len(args)): nodes.append(parse_specifier(args[i], graph, layer_node_map)) layer_name = 'Append_' for node in nodes: layer_name = layer_name + node + "_" if layer_name not in layer_node_map: concat_name = graph.unique_id(prefix=layer_name) graph.add_node(concat_name, parameters=None, op='concat', kind='op') layer_node_map[layer_name] = concat_name i = 0 Node(graph, concat_name).add_sequence_of_ports('in', range(len(nodes))) for node in nodes: out_port = len(Node(graph, node).out_nodes()) Node(graph, node).add_output_port(out_port) graph.create_edge( Node(graph, node), Node(graph, concat_name), out_port, i, create_edge_attrs(node, concat_name, node, i, out_port)) i = i + 1 else: concat_name = layer_node_map[layer_name] return concat_name elif spec == b'Offset': node = parse_specifier(args[0], graph, layer_node_map) t = int(args[1]) if len(args) > 2: raise Error("ModelOptimizer supports only 2 arguments for Offset") layer_name = 'Offset_' + node + '_' if t < 0: layer_name = layer_name + '_' + str(-t) else: layer_name = layer_name + str(t) if layer_name not in layer_node_map: memory_name = graph.unique_id(prefix=layer_name) layer_node_map[layer_name] = memory_name memory_name_2 = memory_name + '_out' graph.add_node(memory_name, parameters=dict(t=t, pair_name=memory_name_2, has_default=False), op='MemoryOffset', kind='op') out_port = len(Node(graph, node).out_nodes()) in_port = len(Node(graph, memory_name).in_nodes()) Node(graph, memory_name).add_input_port(in_port) Node(graph, node).add_output_port(out_port, skip_if_exist=True) graph.create_edge( Node(graph, node), Node(graph, memory_name), out_port, in_port, create_edge_attrs(node, memory_name, node, in_port, out_port)) else: memory_name = layer_node_map[layer_name] return memory_name elif spec == b'Sum': nodes = [] for i in range(len(args)): nodes.append(parse_specifier(args[i], graph, layer_node_map)) layer_name = 'Sum_' for node in nodes: layer_name = layer_name + node + "_" if layer_name not in layer_node_map: sum_name = graph.unique_id(prefix=layer_name) graph.add_node(sum_name, parameters=None, op='Add', kind='op') layer_node_map[layer_name] = sum_name else: sum_name = layer_node_map[layer_name] for i, node in enumerate(nodes): out_port = len(Node(graph, node).out_nodes()) Node(graph, node).add_output_port(out_port, skip_if_exist=True) Node(graph, sum_name).add_input_port(i) graph.add_edge(node, sum_name, **create_edge_attrs(node, sum_name, node, i)) return sum_name elif spec == b'IfDefined': node_id = parse_specifier(args[0], graph, layer_node_map) node = Node(graph, node_id) if node.op == 'MemoryOffset': node['parameters']['has_default'] = True return node_id elif spec == b'ReplaceIndex': node = parse_specifier(args[0], graph, layer_node_map) return node elif spec == b'Scale': node_name = parse_specifier(args[1], graph, layer_node_map) scale_value = float(args[0]) layer_name = '{}/Mul/{}'.format(node_name, scale_value) if layer_name not in layer_node_map: scale_name = graph.unique_id(prefix=layer_name) scale_node = Mul(graph, {'name': scale_name}).create_node() layer_node_map[layer_name] = scale_name scale_const_name = 'Const_{}'.format(scale_value) const_node = Const(graph, { 'name': scale_const_name, 'value': float_array([scale_value]) }).create_node() node = Node(graph, node_name) graph.create_edge( const_node, scale_node, 0, 0, create_edge_attrs(const_node.id, scale_node.id, const_node.id)) out_port = len(node.out_nodes()) graph.create_edge( node, scale_node, out_port, 1, create_edge_attrs(node_name, scale_node.id, node_name, 1, out_port)) else: scale_name = layer_node_map[layer_name] return scale_name
def read_node(file_descr, graph, component_layer_map, layer_node_map): s = file_descr.readline() if s == b'\n': return False tokens = s.split(b' ') if tokens[0] == b'input-node': in_name = s[s.find(b'name=') + len(b'name='):].split(b' ')[0] in_name = str(in_name).strip('b').replace('\'', "") in_shape = np.array( [1, s[s.find(b'dim=') + len(b'dim='):].split(b' ')[0]], dtype=np.int) if in_name not in layer_node_map: graph.add_node(in_name, name=in_name, kind='op', op='Parameter', parameters=None, shape=in_shape) layer_node_map[in_name] = in_name else: Node(graph, in_name)['op'] = 'Parameter' Node(graph, in_name)['shape'] = in_shape elif tokens[0] == b'component-node': layer_name = s[s.find(b'name=') + len(b'name='):].split(b' ')[0] layer_name = str(layer_name).strip('b').replace('\'', "") component_name = s[s.find(b'component=') + len(b'component='):].split(b' ')[0] if layer_name not in layer_node_map: node_name = graph.unique_id(prefix=layer_name) graph.add_node(node_name, parameters=None, op=None, kind='op') layer_node_map[layer_name] = node_name else: node_name = layer_node_map[layer_name] if component_name in component_layer_map: component_layer_map[component_name].append(node_name) else: component_layer_map[component_name] = [node_name] # parse input in_node_id = parse_input_for_node(s[s.find(b'input=') + 6:], graph, layer_node_map) # don't create cyclic edges node to itself to avoid removing later if in_node_id != node_name: out_port = len(Node(graph, in_node_id).out_nodes()) in_port = len(Node(graph, node_name).in_nodes()) Node(graph, node_name).add_input_port(in_port) Node(graph, in_node_id).add_output_port(out_port, skip_if_exist=True) graph.add_edge( in_node_id, node_name, **create_edge_attrs(in_node_id, node_name, in_node_id, in_port, out_port)) elif tokens[0] == b'output-node': layer_name = s[s.find(b'name=') + len(b'name='):].split(b' ')[0] layer_name = str(layer_name).strip('b').replace('\'', "") node_name = graph.unique_id(prefix=layer_name) graph.add_node(node_name, parameters=None, op='Identity', kind='op') out_name = graph.unique_id(prefix=node_name + "_out") graph.add_node(out_name, parameters=None, op='Result', kind='op') Node(graph, node_name).add_input_port(0) Node(graph, node_name).add_output_port(0) Node(graph, out_name).add_input_port(0) graph.add_edge(node_name, out_name, **create_edge_attrs(node_name, out_name, node_name)) # parse input in_node_id = parse_input_for_node( s[s.find(b'input=') + len(b'input='):], graph, layer_node_map) out_port = len(Node(graph, in_node_id).out_nodes()) Node(graph, in_node_id).add_output_port(out_port) graph.create_edge( Node(graph, in_node_id), Node(graph, node_name), out_port, 0, create_edge_attrs(in_node_id, node_name, in_node_id, 0, out_port)) objective_type = s[s.find(b'objective=') + 10:].split(b' ')[0].split(b'\n')[0] if objective_type != b'linear': raise Error( "Unsupported objective-type for output {}".format(node_name)) elif tokens[0] == b'dim-range-node': layer_name = s[s.find(b'name=') + len(b'name='):].split(b' ')[0] layer_name = str(layer_name).strip('b').replace('\'', "") offset = int(s[s.find(b'dim-offset=') + len(b'dim-offset='):].split(b' ')[0]) dim = int(s[s.find(b'dim=') + len(b'dim='):].split(b' ')[0]) if layer_name in layer_node_map: node_name = layer_node_map[layer_name] node = Node(graph, node_name) node['parameters'] = { 'offset': np.array([offset]), 'dim': np.array([dim]), 'axis': np.array([1]) } node['op'] = 'Crop' else: node_name = graph.unique_id(prefix=layer_name) graph.add_node(node_name, parameters={ 'offset': np.array([offset]), 'dim': np.array([dim]), 'axis': np.array([1]) }, op='Crop', kind='op') layer_node_map[layer_name] = node_name node = Node(graph, node_name) in_node_id = parse_input_for_node( s[s.find(b'input-node=') + len(b'input-node='):], graph, layer_node_map) out_port = len(Node(graph, in_node_id).out_nodes()) in_port = len(Node(graph, node_name).in_nodes()) node.add_input_port(in_port) Node(graph, in_node_id).add_output_port(out_port) graph.create_edge( Node(graph, in_node_id), node, out_port, in_port, create_edge_attrs(in_node_id, node_name, in_node_id, in_port, out_port)) # read dim info where possible to simplify shape calculation for MemoryOffset # shape calculation for MemoryOffset can't be done through shape of previous layer because # it is separated in 2 parts to remove cycle from graph for o_n_name, params in node.get_outputs(): o_n = Node(graph, o_n_name) if o_n['op'] == 'MemoryOffset': o_n['parameters']['element_size'] = int64_array([1, dim]) else: raise Error("Unsupported node specifier {}".format(tokens[0])) return True
def load_parallel_component(file_descr, graph: Graph, prev_layer_id): """ Load ParallelComponent of the Kaldi model. ParallelComponent contains parallel nested networks. VariadicSplit is inserted before nested networks. Outputs of nested networks concatenate with layer Concat. :param file_descr: descriptor of the model file :param graph: graph with the topology. :param prev_layer_id: id of the input layers for parallel component layer :return: id of the concat layer - last layer of the parallel component layers """ nnet_count = read_token_value(file_descr, b'<NestedNnetCount>') log.debug( 'Model contains parallel component with {} nested networks'.format( nnet_count)) split_points = [] outputs = [] inputs = [] for i in range(nnet_count): read_token_value(file_descr, b'<NestedNnet>') collect_until_token(file_descr, b'<Nnet>') g = Graph() load_kalid_nnet1_model(g, file_descr, 'Nested_net_{}'.format(i)) # input to nnet1 models is of a rank 1 but we also insert batch_size to 0th axis # 1st axis contains input_size of the nested subnetwork # we split input from the main network to subnetworks input_node = Node(g, 'Parameter') split_points.append(input_node['shape'][1]) g.remove_node(input_node.id) mapping = { node: graph.unique_id(node) for node in g.nodes(data=False) if node in graph } g = nx.relabel_nodes(g, mapping) for val in mapping.values(): g.node[val]['name'] = val graph.add_nodes_from(g.nodes(data=True)) graph.add_edges_from(g.edges(data=True)) sorted_nodes = tuple(nx.topological_sort(g)) outputs.append(Node(graph, sorted_nodes[-1])) inputs.append(Node(graph, sorted_nodes[0])) split_id = graph.unique_id(prefix='NestedNets/VariadicSplit') attrs = { 'out_ports_count': nnet_count, 'size_splits': split_points, 'axis': 1, 'name': split_id } variadic_split_node = AttributedVariadicSplit(graph, attrs).create_node() prev_layer_node = Node(graph, prev_layer_id) prev_layer_node.add_output_port(0) graph.create_edge( prev_layer_node, variadic_split_node, 0, 0, create_edge_attrs(prev_layer_id, variadic_split_node.id, prev_layer_id)) concat_id = graph.unique_id(prefix='Concat') graph.add_node(concat_id, parameters=None, op='concat', kind='op') concat_node = Node(graph, concat_id) # Connect each output of variadic_split_node to each subnetwork's inputs in ParallelComponent # and each subnetwork's output to concat_node for i, (input_node, output_node) in enumerate(zip(inputs, outputs)): output_node.add_output_port(0) concat_node.add_input_port(i) graph.create_edge( output_node, concat_node, 0, i, create_edge_attrs(output_node.id, concat_id, output_node.id, i, 0)) graph.create_edge( variadic_split_node, input_node, i, 0, create_edge_attrs(variadic_split_node.id, input_node.id, variadic_split_node.id, 0, i)) return concat_id
def load_parallel_component(file_descr, graph: Graph, prev_layer_id): """ Load ParallelComponent of the Kaldi model. ParallelComponent contains parallel nested networks. Slice is inserted before nested networks. Outputs of nested networks concatenate with layer Concat. :param file_descr: descriptor of the model file :param graph: graph with the topology. :param prev_layer_id: id of the input layers for parallel component layer :return: id of the concat layer - last layer of the parallel component layers """ nnet_count = read_token_value(file_descr, b'<NestedNnetCount>') log.debug( 'Model contains parallel component with {} nested networks'.format( nnet_count)) slice_id = graph.unique_id(prefix='Slice') graph.add_node(slice_id, parameters=None, op='slice', kind='op') slice_node = Node(graph, slice_id) graph.add_edge(prev_layer_id, slice_id, **create_edge_attrs(prev_layer_id, slice_id)) slices_points = [] outputs = [] for i in range(nnet_count): read_token_value(file_descr, b'<NestedNnet>') collect_until_token(file_descr, b'<Nnet>') g, shape = load_kalid_nnet1_model(file_descr, 'Nested_net_{}'.format(i)) input_nodes = [ n for n in graph.nodes(data=True) if n[1]['op'] == 'Input' ] if i != nnet_count - 1: slices_points.append(shape[1]) g.remove_node(input_nodes[0][0]) mapping = { node: graph.unique_id(node) for node in g.nodes(data=False) if node in graph } g = nx.relabel_nodes(g, mapping) for val in mapping.values(): g.node[val]['name'] = val graph.add_nodes_from(g.nodes(data=True)) graph.add_edges_from(g.edges(data=True)) sorted_nodes = tuple(nx.topological_sort(g)) edge_attrs = create_edge_attrs(slice_id, sorted_nodes[0]) edge_attrs['out'] = i graph.add_edge(slice_id, sorted_nodes[0], **edge_attrs) outputs.append(sorted_nodes[-1]) packed_sp = struct.pack("B", 4) + struct.pack("I", len(slices_points)) for i in slices_points: packed_sp += struct.pack("I", i) slice_node.parameters = io.BytesIO(packed_sp) concat_id = graph.unique_id(prefix='Concat') graph.add_node(concat_id, parameters=None, op='concat', kind='op') for i, output in enumerate(outputs): edge_attrs = create_edge_attrs(output, concat_id) edge_attrs['in'] = i graph.add_edge(output, concat_id, **edge_attrs) return concat_id