|
- from icosagon.data import Data
- from icosagon.bulkdec import BulkDecodeLayer
- from icosagon.input import OneHotInputLayer
- from icosagon.convlayer import DecagonLayer
- import torch
-
-
- def test_bulk_decode_layer_01():
- data = Data()
- data.add_node_type('Dummy', 100)
- fam = data.add_relation_family('Dummy-Dummy', 0, 0, False)
- fam.add_relation_type('Dummy Relation 1',
- torch.rand((100, 100), dtype=torch.float32).round().to_sparse())
-
- in_layer = OneHotInputLayer(data)
- d_layer = DecagonLayer(in_layer.output_dim, 32, data)
- dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
- keep_prob=1., activation=lambda x: x)
- seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)
-
- pred = seq(None)
-
- assert isinstance(pred, list)
- assert len(pred) == len(data.relation_families)
- assert isinstance(pred[0], torch.Tensor)
- assert len(pred[0].shape) == 3
- assert len(pred[0]) == len(data.relation_families[0].relation_types)
- assert pred[0].shape[1] == data.node_types[0].count
- assert pred[0].shape[2] == data.node_types[0].count
-
-
- def test_bulk_decode_layer_02():
- data = Data()
- data.add_node_type('Foo', 100)
- data.add_node_type('Bar', 50)
- fam = data.add_relation_family('Foo-Bar', 0, 1, False)
- fam.add_relation_type('Foobar Relation 1',
- torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
- torch.rand((50, 100), dtype=torch.float32).round().to_sparse())
-
- in_layer = OneHotInputLayer(data)
- d_layer = DecagonLayer(in_layer.output_dim, 32, data)
- dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
- keep_prob=1., activation=lambda x: x)
- seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)
-
- pred = seq(None)
-
- assert isinstance(pred, list)
- assert len(pred) == len(data.relation_families)
- assert isinstance(pred[0], torch.Tensor)
- assert len(pred[0].shape) == 3
- assert len(pred[0]) == len(data.relation_families[0].relation_types)
- assert pred[0].shape[1] == data.node_types[0].count
- assert pred[0].shape[2] == data.node_types[1].count
-
-
- def test_bulk_decode_layer_03():
- data = Data()
- data.add_node_type('Foo', 100)
- data.add_node_type('Bar', 50)
- fam = data.add_relation_family('Foo-Bar', 0, 1, False)
- fam.add_relation_type('Foobar Relation 1',
- torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
- torch.rand((50, 100), dtype=torch.float32).round().to_sparse())
- fam.add_relation_type('Foobar Relation 2',
- torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
- torch.rand((50, 100), dtype=torch.float32).round().to_sparse())
-
- in_layer = OneHotInputLayer(data)
- d_layer = DecagonLayer(in_layer.output_dim, 32, data)
- dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
- keep_prob=1., activation=lambda x: x)
- seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)
-
- pred = seq(None)
-
- assert isinstance(pred, list)
- assert len(pred) == len(data.relation_families)
- assert isinstance(pred[0], torch.Tensor)
- assert len(pred[0].shape) == 3
- assert len(pred[0]) == len(data.relation_families[0].relation_types)
- assert pred[0].shape[1] == data.node_types[0].count
- assert pred[0].shape[2] == data.node_types[1].count
-
-
- def test_bulk_decode_layer_03_big():
- data = Data()
- data.add_node_type('Foo', 2000)
- data.add_node_type('Bar', 2100)
- fam = data.add_relation_family('Foo-Bar', 0, 1, False)
- fam.add_relation_type('Foobar Relation 1',
- torch.rand((2000, 2100), dtype=torch.float32).round().to_sparse(),
- torch.rand((2100, 2000), dtype=torch.float32).round().to_sparse())
- fam.add_relation_type('Foobar Relation 2',
- torch.rand((2000, 2100), dtype=torch.float32).round().to_sparse(),
- torch.rand((2100, 2000), dtype=torch.float32).round().to_sparse())
-
- in_layer = OneHotInputLayer(data)
- d_layer = DecagonLayer(in_layer.output_dim, 32, data)
- dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
- keep_prob=1., activation=lambda x: x)
- seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)
-
- pred = seq(None)
-
- assert isinstance(pred, list)
- assert len(pred) == len(data.relation_families)
- assert isinstance(pred[0], torch.Tensor)
- assert len(pred[0].shape) == 3
- assert len(pred[0]) == len(data.relation_families[0].relation_types)
- assert pred[0].shape[1] == data.node_types[0].count
- assert pred[0].shape[2] == data.node_types[1].count
|