From cef1d1110ab1e24bf8830c3fdcedc3a029769463 Mon Sep 17 00:00:00 2001 From: Stanislaw Adaszewski Date: Tue, 21 Jul 2020 12:51:51 +0200 Subject: [PATCH] Add more tests in test_bulkdec. --- tests/icosagon/test_bulkdec.py | 127 +++++++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) diff --git a/tests/icosagon/test_bulkdec.py b/tests/icosagon/test_bulkdec.py index c1e4670..b4ed3d6 100644 --- a/tests/icosagon/test_bulkdec.py +++ b/tests/icosagon/test_bulkdec.py @@ -3,6 +3,9 @@ from icosagon.bulkdec import BulkDecodeLayer from icosagon.input import OneHotInputLayer from icosagon.convlayer import DecagonLayer import torch +import pytest +import time +import sys def test_bulk_decode_layer_01(): @@ -111,3 +114,127 @@ def test_bulk_decode_layer_03_big(): 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_huge_gpu(): + if torch.cuda.device_count() == 0: + pytest.skip('test_bulk_decode_layer_03_huge_gpu() requires CUDA support') + + device = torch.device('cuda:0') + data = Data() + data.add_node_type('Foo', 20000) + data.add_node_type('Bar', 21000) + fam = data.add_relation_family('Foo-Bar', 0, 1, False) + print('Adding Foobar Relation 1...') + fam.add_relation_type('Foobar Relation 1', + torch.rand((20000, 21000), dtype=torch.float32).round().to_sparse().to(device), + torch.rand((21000, 20000), dtype=torch.float32).round().to_sparse().to(device)) + print('Adding Foobar Relation 2...') + fam.add_relation_type('Foobar Relation 2', + torch.rand((20000, 21000), dtype=torch.float32).round().to_sparse().to(device), + torch.rand((21000, 20000), dtype=torch.float32).round().to_sparse().to(device)) + + 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) + seq = seq.to(device) + + print('Starting forward pass...') + t = time.time() + pred = seq(None) + print('Elapsed:', time.time() - t) + + 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_04_huge_multirel_gpu(): + if torch.cuda.device_count() == 0: + pytest.skip('test_bulk_decode_layer_04_huge_multirel_gpu() requires CUDA support') + + if torch.cuda.get_device_properties(0).total_memory < 64000000000: + pytest.skip('test_bulk_decode_layer_04_huge_multirel_gpu() requires GPU with 64GB of memory') + + device = torch.device('cuda:0') + data = Data() + data.add_node_type('Foo', 20000) + data.add_node_type('Bar', 21000) + fam = data.add_relation_family('Foo-Bar', 0, 1, False) + print('Generating adj_mat ...') + adj_mat = torch.rand((20000, 21000), dtype=torch.float32).round().to_sparse().to(device) + print('Generating adj_mat_back ...') + adj_mat_back = torch.rand((21000, 20000), dtype=torch.float32).round().to_sparse().to(device) + print('Adding relations ...') + for i in range(1300): + sys.stdout.write('.') + sys.stdout.flush() + fam.add_relation_type(f'Foobar Relation {i}', adj_mat, adj_mat_back) + print() + + 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) + seq = seq.to(device) + + print('Starting forward pass...') + t = time.time() + pred = seq(None) + print('Elapsed:', time.time() - t) + + 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_04_big_multirel_gpu(): + if torch.cuda.device_count() == 0: + pytest.skip('test_bulk_decode_layer_04_big_multirel_gpu() requires CUDA support') + + device = torch.device('cuda:0') + data = Data() + data.add_node_type('Foo', 2000) + data.add_node_type('Bar', 2100) + fam = data.add_relation_family('Foo-Bar', 0, 1, False) + print('Generating adj_mat ...') + adj_mat = torch.rand((2000, 2100), dtype=torch.float32).round().to_sparse().to(device) + print('Generating adj_mat_back ...') + adj_mat_back = torch.rand((2100, 2000), dtype=torch.float32).round().to_sparse().to(device) + print('Adding relations ...') + for i in range(1300): + sys.stdout.write('.') + sys.stdout.flush() + fam.add_relation_type(f'Foobar Relation {i}', adj_mat, adj_mat_back) + print() + + 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) + seq = seq.to(device) + + print('Starting forward pass...') + t = time.time() + pred = seq(None) + print('Elapsed:', time.time() - t) + + 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