from icosagon.data import Data from icosagon.trainprep import prepare_training, \ TrainValTest from icosagon.model import Model from icosagon.trainloop import TrainLoop import torch import pytest import pdb import time def test_train_loop_01(): d = Data() d.add_node_type('Dummy', 10) fam = d.add_relation_family('Dummy-Dummy', 0, 0, False) fam.add_relation_type('Dummy Rel', torch.rand(10, 10).round()) prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) m = Model(prep_d) loop = TrainLoop(m) assert loop.model == m assert loop.lr == 0.001 assert loop.loss == torch.nn.functional.binary_cross_entropy_with_logits assert loop.batch_size == 100 def test_train_loop_02(): d = Data() d.add_node_type('Dummy', 10) fam = d.add_relation_family('Dummy-Dummy', 0, 0, False) fam.add_relation_type('Dummy Rel', torch.rand(10, 10).round()) prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) m = Model(prep_d) for prm in m.parameters(): print(prm.shape, prm.is_leaf, prm.requires_grad) loop = TrainLoop(m) loop.run_epoch() for prm in m.parameters(): print(prm.shape, prm.is_leaf, prm.requires_grad) def test_train_loop_03(): # pdb.set_trace() if torch.cuda.device_count() == 0: pytest.skip('CUDA required for this test') adj_mat = torch.rand(10, 10).round() dev = torch.device('cuda:0') adj_mat = adj_mat.to(dev) d = Data() d.add_node_type('Dummy', 10) fam = d.add_relation_family('Dummy-Dummy', 0, 0, False) fam.add_relation_type('Dummy Rel', adj_mat) prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) # pdb.set_trace() m = Model(prep_d) m = m.to(dev) print(list(m.parameters())) for prm in m.parameters(): assert prm.device == dev loop = TrainLoop(m) loop.run_epoch() def test_timing_01(): adj_mat = (torch.rand(2000, 2000) < .001).to(torch.float32).to_sparse() rep = torch.eye(2000).requires_grad_(True) t = time.time() for _ in range(1300): _ = torch.sparse.mm(adj_mat, rep) print('Elapsed:', time.time() - t) def test_timing_02(): adj_mat = (torch.rand(2000, 2000) < .001).to(torch.float32) adj_mat_batch = [adj_mat.view(1, 2000, 2000)] * 1300 adj_mat_batch = torch.cat(adj_mat_batch) rep = torch.eye(2000).requires_grad_(True) t = time.time() res = torch.matmul(adj_mat_batch, rep) print('Elapsed:', time.time() - t) print('res.shape:', res.shape) def test_timing_03(): adj_mat = (torch.rand(2000, 2000) < .001).to(torch.float32) adj_mat_batch = [adj_mat.view(1, 2000, 2000).to_sparse()] * 1300 adj_mat_batch = torch.cat(adj_mat_batch) rep = torch.eye(2000).requires_grad_(True) rep_batch = [rep.view(1, 2000, 2000)] * 1300 rep_batch = torch.cat(rep_batch) t = time.time() with pytest.raises(RuntimeError): _ = torch.bmm(adj_mat_batch, rep) print('Elapsed:', time.time() - t) def test_timing_04(): adj_mat = (torch.rand(2000, 2000) < .0001).to(torch.float32).to_sparse() rep = torch.eye(2000).requires_grad_(True) t = time.time() for _ in range(1300): _ = torch.sparse.mm(adj_mat, rep) print('Elapsed:', time.time() - t) def test_timing_05(): if torch.cuda.device_count() == 0: pytest.skip('Test requires CUDA') dev = torch.device('cuda:0') adj_mat = (torch.rand(2000, 2000) < .001).to(torch.float32).to_sparse().to(dev) rep = torch.eye(2000).requires_grad_(True).to(dev) t = time.time() for _ in range(1300): _ = torch.sparse.mm(adj_mat, rep) torch.cuda.synchronize() print('Elapsed:', time.time() - t)