Optional dependencies (#4411)

* initial commit * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update * withPackage decorator * update * install test * update * typo * update * update * fix hydra test * update * update * update * update * fix test * update * update * update Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
pyg-team · Apr 5, 2022 · 8281be4 · 8281be4
1 parent bfe9ce1
commit 8281be4
Show file tree

Hide file tree

Showing 41 changed files with 306 additions and 127 deletions.
diff --git a/.github/workflows/install.yml b/.github/workflows/install.yml
@@ -35,11 +35,15 @@ jobs:
 
       - name: Install main package
         run: |
-          pip install .
+          pip install -e .[test]
 
       - name: Test imports
         run: |
           python -c "import torch_geometric"
           python -c "import torch_geometric.datasets"
           python -c "import torch_geometric.nn"
           python -c "import torch_geometric.graphgym"
+
+      - name: Run tests
+        run: |
+          pytest
diff --git a/conda/pyg/meta.yaml b/conda/pyg/meta.yaml
@@ -17,17 +17,13 @@ requirements:
     - pytorch-scatter
     - pytorch-sparse
     - pytorch-cluster
-    - networkx>=2.4
     - tqdm
     - jinja2
     - pyparsing
-    - yacs
-    - PyYAML
-    - scikit-learn
-    - python-louvain
-    - pandas
-    - requests
     - numpy
+    - scipy
+    - requests
+    - scikit-learn
 
 build:
   string: py{{ environ.get('PYTHON_VERSION').replace('.', '') }}_torch_{{ environ['TORCH_VERSION'] }}_{{ environ['CUDA_VERSION'] }}

diff --git a/conda/pytorch-geometric/meta.yaml b/conda/pytorch-geometric/meta.yaml
@@ -17,17 +17,13 @@ requirements:
     - pytorch-scatter
     - pytorch-sparse
     - pytorch-cluster
-    - networkx>=2.4
     - tqdm
     - jinja2
     - pyparsing
-    - yacs
-    - PyYAML
-    - scikit-learn
-    - python-louvain
-    - pandas
-    - requests
     - numpy
+    - scipy
+    - requests
+    - scikit-learn
 
 build:
   string: py{{ environ.get('PYTHON_VERSION').replace('.', '') }}_torch_{{ environ['TORCH_VERSION'] }}_{{ environ['CUDA_VERSION'] }}

diff --git a/docs/README.md b/docs/README.md
@@ -5,10 +5,9 @@ To build the documentation:
 1. [Build and install](https://github.com/pyg-team/pytorch_geometric/blob/master/CONTRIBUTING.md#developing-pytorch-geometric) PyG from source.
 2. Install [Sphinx](https://www.sphinx-doc.org/en/master/) via `pip install sphinx sphinx_rtd_theme`.
 3. Generate the documentation file via:
-
-```
-cd docs
-make html
-```
+   ```
+   cd docs
+   make html
+   ```
 
 The documentation is now available to view by opening `docs/build/html/index.html`.
diff --git a/setup.py b/setup.py
@@ -7,7 +7,6 @@
     'tqdm',
     'numpy',
     'scipy',
-    'pandas',
     'jinja2',
     'requests',
     'pyparsing',
@@ -18,6 +17,7 @@
     'h5py',
     'yacs',
     'numba',
+    'pandas',
     'captum',
     'rdflib',
     'trimesh',
@@ -28,7 +28,6 @@
     'scikit-image',
     'pytorch-memlab',
     'torchmetrics>=0.7',
-    'class-resolver>=0.3.2',
 ]
 
 test_requires = [

diff --git a/test/data/test_lightning_datamodule.py b/test/data/test_lightning_datamodule.py
@@ -6,14 +6,12 @@
 
 from torch_geometric.data import LightningDataset, LightningNodeData
 from torch_geometric.nn import global_mean_pool
-from torch_geometric.testing import onlyFullTest
+from torch_geometric.testing import onlyFullTest, withCUDA, withPackage
 
 try:
     from pytorch_lightning import LightningModule
-    no_pytorch_lightning = False
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     LightningModule = torch.nn.Module
-    no_pytorch_lightning = True
 
 
 class LinearGraphModule(LightningModule):
@@ -53,9 +51,9 @@ def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=0.01)
 
 
+@withCUDA
 @onlyFullTest
-@pytest.mark.skipif(no_pytorch_lightning, reason='PL not available')
-@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA not available')
+@withPackage('pytorch_lightning')
 @pytest.mark.parametrize('strategy_type', [None, 'ddp_spawn'])
 def test_lightning_dataset(get_dataset, strategy_type):
     import pytorch_lightning as pl
@@ -144,9 +142,9 @@ def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=0.01)
 
 
+@withCUDA
 @onlyFullTest
-@pytest.mark.skipif(no_pytorch_lightning, reason='PL not available')
-@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA not available')
+@withPackage('pytorch_lightning')
 @pytest.mark.parametrize('loader', ['full', 'neighbor'])
 @pytest.mark.parametrize('strategy_type', [None, 'ddp_spawn'])
 def test_lightning_node_data(get_dataset, strategy_type, loader):
@@ -239,9 +237,9 @@ def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=0.01)
 
 
+@withCUDA
 @onlyFullTest
-@pytest.mark.skipif(no_pytorch_lightning, reason='PL not available')
-@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA not available')
+@withPackage('pytorch_lightning')
 def test_lightning_hetero_node_data(get_dataset):
     import pytorch_lightning as pl
 

diff --git a/test/graphgym/test_config_store.py b/test/graphgym/test_config_store.py
@@ -1,10 +1,13 @@
-import hydra
-from omegaconf import DictConfig
-
-from torch_geometric.graphgym.config_store import fill_config_store
+from torch_geometric.testing import withPackage
 
 
+@withPackage('hydra')
 def test_config_store():
+    import hydra
+    from omegaconf import DictConfig
+
+    from torch_geometric.graphgym.config_store import fill_config_store
+
     fill_config_store()
 
     with hydra.initialize(config_path='.'):

diff --git a/test/graphgym/test_graphgym.py b/test/graphgym/test_graphgym.py
@@ -29,6 +29,7 @@
     auto_select_device,
     params_count,
 )
+from torch_geometric.testing import withPackage
 
 num_trivial_metric_calls = 0
 
@@ -39,6 +40,7 @@ def trivial_metric(true, pred, task_type):
     return 1
 
 
+@withPackage('yacs')
 @pytest.mark.parametrize('auto_resume', [True, False])
 @pytest.mark.parametrize('skip_train_eval', [True, False])
 @pytest.mark.parametrize('use_trivial_metric', [True, False])

diff --git a/test/graphgym/test_register.py b/test/graphgym/test_register.py
@@ -1,13 +1,15 @@
 import torch
 
 import torch_geometric.graphgym.register as register
+from torch_geometric.testing import withPackage
 
 
 @register.register_act('identity')
 def identity_act(x: torch.Tensor) -> torch.Tensor:
     return x
 
 
+@withPackage('yacs')
 def test_register():
     assert len(register.act_dict) == 8
     assert list(register.act_dict.keys()) == [

diff --git a/test/nn/conv/test_edge_conv.py b/test/nn/conv/test_edge_conv.py
@@ -5,7 +5,7 @@
 from torch_sparse import SparseTensor
 
 from torch_geometric.nn import DynamicEdgeConv, EdgeConv
-from torch_geometric.testing import is_full_test
+from torch_geometric.testing import is_full_test, withPackage
 
 
 def test_edge_conv_conv():
@@ -56,6 +56,7 @@ def test_edge_conv_conv():
         assert jit((x1, x2), adj.t()).tolist() == out2.tolist()
 
 
+@withPackage('torch_cluster')
 def test_dynamic_edge_conv_conv():
     x1 = torch.randn(8, 16)
     x2 = torch.randn(4, 16)

diff --git a/test/nn/conv/test_gravnet_conv.py b/test/nn/conv/test_gravnet_conv.py
@@ -1,9 +1,10 @@
 import torch
 
 from torch_geometric.nn import GravNetConv
-from torch_geometric.testing import is_full_test
+from torch_geometric.testing import is_full_test, withPackage
 
 
+@withPackage('torch_cluster')
 def test_gravnet_conv():
     x1 = torch.randn(8, 16)
     x2 = torch.randn(4, 16)

diff --git a/test/nn/conv/test_message_passing.py b/test/nn/conv/test_message_passing.py
@@ -65,6 +65,32 @@ def test_my_conv():
     assert conv(x1, adj.t()).tolist() == out.tolist()
     conv.fuse = True
 
+    adj = adj.sparse_resize((4, 2))
+    conv = MyConv((8, 16), 32)
+    out1 = conv((x1, x2), edge_index, value)
+    out2 = conv((x1, None), edge_index, value, (4, 2))
+    assert out1.size() == (2, 32)
+    assert out2.size() == (2, 32)
+    assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
+    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
+    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
+    conv.fuse = False
+    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
+    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
+    conv.fuse = True
+
+
+def test_my_conv_jittable():
+    x1 = torch.randn(4, 8)
+    x2 = torch.randn(2, 16)
+    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+    row, col = edge_index
+    value = torch.randn(row.size(0))
+    adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
+
+    conv = MyConv(8, 32)
+    out = conv(x1, edge_index, value)
+
     t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
     assert jit(x1, edge_index, value).tolist() == out.tolist()
@@ -81,15 +107,6 @@ def test_my_conv():
     conv = MyConv((8, 16), 32)
     out1 = conv((x1, x2), edge_index, value)
     out2 = conv((x1, None), edge_index, value, (4, 2))
-    assert out1.size() == (2, 32)
-    assert out2.size() == (2, 32)
-    assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
-    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
-    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
-    conv.fuse = False
-    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
-    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
-    conv.fuse = True
 
     t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
@@ -166,6 +183,16 @@ def test_my_multiple_aggr_conv():
     assert not torch.allclose(out[:, 16:32], out[:, 32:48])
     assert torch.allclose(conv(x, adj.t()), out)
 
+
+def test_my_multiple_aggr_conv_jittable():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+    row, col = edge_index
+    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
+
+    conv = MyMultipleAggrConv()
+    out = conv(x, edge_index)
+
     t = '(Tensor, Tensor) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
     assert torch.allclose(jit(x, edge_index), out)
@@ -226,13 +253,23 @@ def test_my_edge_conv():
     assert torch.allclose(out, expected)
     assert torch.allclose(conv(x, adj.t()), out)
 
+
+def test_my_edge_conv_jittable():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+    row, col = edge_index
+    adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
+
+    conv = MyEdgeConv()
+    out = conv(x, edge_index)
+
     t = '(Tensor, Tensor) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit(x, edge_index), expected)
+    assert torch.allclose(jit(x, edge_index), out)
 
     t = '(Tensor, SparseTensor) -> Tensor'
     jit = torch.jit.script(conv.jittable(t))
-    assert torch.allclose(jit(x, adj.t()), expected)
+    assert torch.allclose(jit(x, adj.t()), out)
 
 
 num_pre_hook_calls = 0
@@ -373,8 +410,11 @@ def test_my_default_arg_conv():
     assert conv(x, edge_index).view(-1).tolist() == [0, 0, 0, 0]
     assert conv(x, adj.t()).view(-1).tolist() == [0, 0, 0, 0]
 
-    # This should not succeed in JIT mode.
-    with pytest.raises(RuntimeError):
+
+def test_my_default_arg_conv_jittable():
+    conv = MyDefaultArgConv()
+
+    with pytest.raises(RuntimeError):  # This should not succeed in JIT mode.
         torch.jit.script(conv.jittable())
 
 
@@ -407,6 +447,16 @@ def test_tuple_output():
     out1 = conv(x, edge_index)
     assert isinstance(out1, tuple) and len(out1) == 2
 
+
+def test_tuple_output_jittable():
+    conv = MyMultipleOutputConv()
+
+    x = torch.randn(4, 8)
+    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+
+    out1 = conv(x, edge_index)
+    assert isinstance(out1, tuple) and len(out1) == 2
+
     jit = torch.jit.script(conv.jittable())
     out2 = jit(x, edge_index)
     assert isinstance(out2, tuple) and len(out2) == 2

diff --git a/test/nn/conv/test_spline_conv.py b/test/nn/conv/test_spline_conv.py
@@ -2,9 +2,10 @@
 from torch_sparse import SparseTensor
 
 from torch_geometric.nn import SplineConv
-from torch_geometric.testing import is_full_test
+from torch_geometric.testing import is_full_test, withPackage
 
 
+@withPackage('torch_spline_conv')
 def test_spline_conv():
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)
@@ -56,6 +57,7 @@ def test_spline_conv():
         assert torch.allclose(jit((x1, None), adj.t()), out2)
 
 
+@withPackage('torch_spline_conv')
 def test_lazy_spline_conv():
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)

diff --git a/test/nn/conv/test_x_conv.py b/test/nn/conv/test_x_conv.py
@@ -1,9 +1,10 @@
 import torch
 
 from torch_geometric.nn import XConv
-from torch_geometric.testing import is_full_test
+from torch_geometric.testing import is_full_test, withPackage
 
 
+@withPackage('torch_cluster')
 def test_x_conv():
     x = torch.randn(8, 16)
     pos = torch.rand(8, 3)