def test(self):
graph = build_graph(nodes_attributes, [
('weights_inp', 'aten', {
'in': 0,
'out': 0
}),
('indices_inp', 'aten', {
'in': 1,
'out': 0
}),
('offsets_inp', 'aten', {
'in': 2,
'out': 0
}),
('aten', 'result', {
'in': 0,
'out': 0
}),
], {},
nodes_with_edges_only=True)
graph_ref = build_graph(nodes_attributes, [
('weights_inp', 'emb_bag', {
'in': 0,
'out': 0
}),
('indices_inp', 'emb_bag', {
'in': 1,
'out': 0
}),
('offsets_inp', 'emb_bag', {
'in': 2,
'out': 0
}),
('emb_bag', 'result', {
'in': 0,
'out': 0
}),
], {},
nodes_with_edges_only=True)
graph.graph['layout'] = 'NCHW'
graph.stage = 'front'
replacer = AtenToEmbeddingBag()
replacer.find_and_replace_pattern(graph)
(flag, resp) = compare_graphs(graph,
graph_ref,
'result',
check_op_attrs=True)
self.assertTrue(flag, resp)
self.assertTrue(graph.node[graph.get_nodes_with_attributes(
op='EmbeddingBag')[0]]['name'] == 'my_aten')
def test(self):
nodes = {
**const('weights_inp', np.random.randn(100, 2)),
**regular_op('indices_inp', {'type': 'Parameter'}),
**regular_op('offsets_inp', {'type': 'Parameter'}),
**regular_op(
'aten', {
'type': None,
'kind': 'op',
'op': 'ATen',
'operator': 'embedding_bag',
'mode': 0,
'name': 'my_aten'
}),
**regular_op(
'emb_bag', {
'type': 'EmbeddingBagOffsetsSum',
'kind': 'op',
'op': 'EmbeddingBagOffsetsSum'
}),
**result('result'),
}
edges = [
('weights_inp', 'aten'),
('indices_inp', 'aten'),
('offsets_inp', 'aten'),
('aten', 'result'),
]
graph = build_graph(nodes, edges)
graph.graph['layout'] = 'NCHW'
graph.stage = 'front'
edges_ref = [
('weights_inp', 'emb_bag'),
('indices_inp', 'emb_bag'),
('offsets_inp', 'emb_bag'),
('emb_bag', 'result'),
]
graph_ref = build_graph(nodes, edges_ref)
AtenToEmbeddingBag().find_and_replace_pattern(graph)
(flag, resp) = compare_graphs(graph, graph_ref, 'result')
self.assertTrue(flag, resp)
def test_per_sample_weights(self):
nodes = {
**const('weights_inp', np.random.randn(100, 2)),
**regular_op('indices_inp', {'type': 'Parameter'}),
**regular_op('offsets_inp', {'type': 'Parameter'}),
**regular_op('per_sample_weights', {'type': 'Parameter'}),
**regular_op(
'aten', {
'type': None,
'kind': 'op',
'op': 'ATen',
'operator': 'embedding_bag',
'mode': 0,
'name': 'my_aten'
}),
**regular_op(
'emb_bag', {
'type': 'EmbeddingBagOffsetsSum',
'kind': 'op',
'op': 'EmbeddingBagOffsetsSum'
}),
**regular_op('WeightsRank', {
'type': None,
'kind': 'op',
'op': 'Rank'
}),
**regular_op('WeightsRank/axis', {
'type': 'Add',
'kind': 'op',
'op': 'Add'
}),
**regular_op('gather1', {
'type': 'Gather',
'kind': 'op',
'op': 'Gather'
}),
**regular_op('gather2', {
'type': 'Gather',
'kind': 'op',
'op': 'Gather'
}),
**regular_op('WeightsShape', {
'type': 'ShapeOf',
'kind': 'op',
'op': 'ShapeOf'
}),
**regular_op('Broadcast', {
'type': 'Broadcast',
'kind': 'op',
'op': 'Broadcast'
}),
**regular_op('Unsqueeze', {
'type': 'Unsqueeze',
'kind': 'op',
'op': 'Unsqueeze'
}),
**const('WeightsShape/Axis', int64_array(0)),
**const('zero1', int64_array(0)),
**const('zero2', int64_array(0)),
**const('Unsqueeze/value', int64_array(0)),
**const('Broadcast/value', int64_array(0)),
**const('neg', int64_array(-1)),
**regular_op('Concat', {
'type': 'Concat',
'kind': 'op',
'op': 'Concat'
}),
**result('result'),
}
edges = [
('weights_inp', 'aten'),
('indices_inp', 'aten'),
('offsets_inp', 'aten'),
('per_sample_weights', 'aten'),
('aten', 'result'),
]
graph = build_graph(nodes, edges, nodes_with_edges_only=True)
graph.graph['layout'] = 'NCHW'
graph.stage = 'front'
edges_ref = [
('weights_inp', 'Concat', {
'in': 0,
'out': 0
}),
('weights_inp', 'WeightsShape', {
'in': 0,
'out': 0
}),
('weights_inp', 'WeightsRank', {
'in': 0,
'out': 0
}),
('WeightsRank', 'WeightsRank/axis'),
('neg', 'WeightsRank/axis'),
('WeightsShape', 'gather1', {
'in': 0,
'out': 0
}),
('WeightsRank/axis', 'gather1'),
('WeightsShape/Axis', 'gather1'),
('WeightsShape', 'gather2', {
'in': 0,
'out': 0
}),
('zero1', 'gather2'),
('zero2', 'gather2'),
('Broadcast/value', 'Broadcast'),
('gather1', 'Broadcast'),
('Broadcast', 'Unsqueeze'),
('Unsqueeze/value', 'Unsqueeze'),
('Unsqueeze', 'Concat'),
('Concat', 'emb_bag'),
('indices_inp', 'emb_bag'),
('offsets_inp', 'emb_bag'),
('gather2', 'emb_bag'),
('per_sample_weights', 'emb_bag'),
('emb_bag', 'result'),
]
graph_ref = build_graph(nodes, edges_ref, nodes_with_edges_only=True)
AtenToEmbeddingBag().find_and_replace_pattern(graph)
(flag, resp) = compare_graphs(graph, graph_ref, 'result')
self.assertTrue(flag, resp)