metaopt · XuehaiPan · Aug 24, 2022 · Aug 8, 2022 · Aug 8, 2022 · Aug 8, 2022
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -29,7 +29,7 @@ concurrency:
 jobs:
   test:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 60
     steps:
       - name: Checkout
         uses: actions/checkout@v3

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,10 +13,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Add question/help/support issue template by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#43](https://github.com/metaopt/TorchOpt/pull/43).
-- Add parallel training on one GPU using functorch.vmap example by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#32](https://github.com/metaopt/TorchOpt/pull/32).
+- Refactor tests using `pytest.mark.parametrize` and enabling parallel testing by [@XuehaiPan](https://github.com/XuehaiPan) in [#55](https://github.com/metaopt/TorchOpt/pull/55).
 - Add maml-omniglot few-shot classification example using functorch.vmap by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#39](https://github.com/metaopt/TorchOpt/pull/39).
-
+- Add parallel training on one GPU using functorch.vmap example by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#32](https://github.com/metaopt/TorchOpt/pull/32).
+- Add question/help/support issue template by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#43](https://github.com/metaopt/TorchOpt/pull/43).
 
 ### Changed
 
@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix RMSProp optimizer by [@XuehaiPan](https://github.com/XuehaiPan) in [#55](https://github.com/metaopt/TorchOpt/pull/55).
 - Fix momentum tracing by [@XuehaiPan](https://github.com/XuehaiPan) in [#58](https://github.com/metaopt/TorchOpt/pull/58).
 - Fix CUDA build for accelerated OP by [@XuehaiPan](https://github.com/XuehaiPan) in [#53](https://github.com/metaopt/TorchOpt/pull/53).
 - Fix gamma error in MAML-RL implementation by [@Benjamin-eecs](https://github.com/Benjamin-eecs) [#47](https://github.com/metaopt/TorchOpt/pull/47).

diff --git a/Makefile b/Makefile
@@ -83,8 +83,9 @@ addlicense-install: go-install
 
 pytest: pytest-install
 	cd tests && \
-	$(PYTHON) -m pytest unit --verbose --color=yes --durations=0 \
-		--cov="$(PROJECT_PATH)" --cov-report=xml --cov-report=term-missing
+	$(PYTHON) -m pytest --verbose --color=yes --durations=0 \
+		--cov="$(PROJECT_NAME)" --cov-report=xml --cov-report=term-missing \
+		.
 
 test: pytest
 

diff --git a/docs/conda-recipe.yaml b/docs/conda-recipe.yaml
@@ -31,7 +31,6 @@ dependencies:
 
   # Learning
   - pytorch::pytorch >= 1.12
-  - pytorch::torchvision
   - pytorch::pytorch-mutex = *=*cpu*
   - pip:
       - functorch >= 0.2

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -1,6 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch >= 1.12
-torchvision
 functorch >= 0.2
 
 --requirement ../requirements.txt

diff --git a/pyproject.toml b/pyproject.toml
@@ -73,7 +73,6 @@ lint = [
     "pre-commit",
 ]
 test = [
-    'torchvision',
     'functorch >= 0.2',
     'pytest',
     'pytest-cov',
@@ -123,10 +122,10 @@ test-command = """
     TORCH_VERSION="$(python -c 'print(__import__("torch").__version__.partition("+")[0])')"
     TEST_TORCH_SPECS="${TEST_TORCH_SPECS:-"${DEFAULT_TEST_TORCH_SPECS}"}"
     for spec in ${TEST_TORCH_SPECS}; do
-        python -m pip uninstall -y torch torchvision
+        python -m pip uninstall -y torch
         export PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/${spec}"
         echo "PIP_EXTRA_INDEX_URL='${PIP_EXTRA_INDEX_URL}'"
-        python -m pip install "torch==${TORCH_VERSION}" torchvision
+        python -m pip install "torch==${TORCH_VERSION}"
         echo "ls ${TORCH_LIB_PATH}"; ls -lh "${TORCH_LIB_PATH}"
         find "${SITE_PACKAGES}/torchopt" -name "*.so" -print0 |
             xargs -0 -I '{}' bash -c "echo 'ldd {}'; ldd '{}'; echo 'patchelf --print-rpath {}'; patchelf --print-rpath '{}'"

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -0,0 +1,20 @@
+# Copyright 2022 MetaOPT Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+
+
+os.environ['PYTHONHASHSEED'] = '0'
+os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -0,0 +1,170 @@
+# Copyright 2022 MetaOPT Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import copy
+import itertools
+import os
+import random
+from typing import Iterable, Optional, Tuple, Union
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+from torch.utils import data
+
+
+BATCH_SIZE = 4
+NUM_UPDATES = 3
+
+MODEL_NUM_INPUTS = 28 * 28  # MNIST
+MODEL_NUM_CLASSES = 10
+MODEL_HIDDEN_SIZE = 64
+
+
+def parametrize(**argvalues) -> pytest.mark.parametrize:
+    arguments = list(argvalues)
+
+    if 'dtype' in argvalues:
+        dtypes = argvalues['dtype']
+        argvalues['dtype'] = dtypes[:1]
+        arguments.remove('dtype')
+        arguments.insert(0, 'dtype')
+
+        argvalues = list(itertools.product(*tuple(map(argvalues.get, arguments))))
+        first_product = argvalues[0]
+        argvalues.extend((dtype,) + first_product[1:] for dtype in dtypes[1:])
+
+    ids = tuple(
+        '-'.join(f'{arg}({val})' for arg, val in zip(arguments, values)) for values in argvalues
+    )
+
+    return pytest.mark.parametrize(arguments, argvalues, ids=ids)
+
+
+def seed_everything(seed: int) -> None:
+    os.environ['PYTHONHASHSEED'] = str(seed)
+
+    random.seed(seed)
+    np.random.seed(seed)
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    try:
+        torch.use_deterministic_algorithms(True)
+    except AttributeError:
+        pass
+
+
+@torch.no_grad()
+def get_models(
+    device: Optional[Union[str, torch.device]] = None, dtype: torch.dtype = torch.float32
+) -> Tuple[nn.Module, nn.Module, nn.Module, data.DataLoader]:
+    seed_everything(seed=42)
+
+    model_base = nn.Sequential(
+        nn.Linear(
+            in_features=MODEL_NUM_INPUTS,
+            out_features=MODEL_HIDDEN_SIZE,
+            bias=True,
+            dtype=dtype,
+        ),
+        nn.ReLU(),
+        nn.Linear(
+            in_features=MODEL_HIDDEN_SIZE,
+            out_features=MODEL_HIDDEN_SIZE,
+            bias=True,
+            dtype=dtype,
+        ),
+        nn.ReLU(),
+        nn.Linear(
+            in_features=MODEL_HIDDEN_SIZE,
+            out_features=MODEL_NUM_CLASSES,
+            bias=True,
+            dtype=dtype,
+        ),
+        nn.Softmax(dim=-1),
+    )
+    for name, param in model_base.named_parameters(recurse=True):
+        if name.endswith('weight'):
+            nn.init.orthogonal_(param)
+        if name.endswith('bias'):
+            param.data.normal_(0, 0.1)
+
+    model = copy.deepcopy(model_base)
+    model_ref = copy.deepcopy(model_base)
+    if device is not None:
+        model_base = model_base.to(device=torch.device(device))
+        model = model.to(device=torch.device(device))
+        model_ref = model_ref.to(device=torch.device(device))
+
+    dataset = data.TensorDataset(
+        torch.randint(0, 1, (BATCH_SIZE * NUM_UPDATES, MODEL_NUM_INPUTS)),
+        torch.randint(0, MODEL_NUM_CLASSES, (BATCH_SIZE * NUM_UPDATES,)),
+    )
+    loader = data.DataLoader(dataset, BATCH_SIZE, shuffle=False)
+
+    return model, model_ref, model_base, loader
+
+
+@torch.no_grad()
+def assert_model_all_close(
+    model: Union[nn.Module, Tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]],
+    model_ref: nn.Module,
+    model_base: nn.Module,
+    dtype: torch.dtype = torch.float32,
+    rtol: Optional[float] = None,
+    atol: Optional[float] = None,
+    equal_nan: bool = False,
+):
+
+    if isinstance(model, tuple):
+        params, buffers = model
+    elif isinstance(model, nn.Module):
+        params = model.parameters()
+        buffers = model.buffers()
+
+    for p, p_ref, p_base in zip(params, model_ref.parameters(), model_base.parameters()):
+        assert_all_close(p, p_ref, base=p_base, rtol=rtol, atol=atol, equal_nan=equal_nan)
+    for b, b_ref, b_base in zip(buffers, model_ref.buffers(), model_base.buffers()):
+        b = b.to(dtype=dtype) if not b.is_floating_point() else b
+        b_ref = b_ref.to(dtype=dtype) if not b_ref.is_floating_point() else b_ref
+        b_base = b_base.to(dtype=dtype) if not b_base.is_floating_point() else b_base
+        assert_all_close(b, b_ref, base=b_base, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+
+@torch.no_grad()
+def assert_all_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    base: torch.Tensor = None,
+    rtol: Optional[float] = None,
+    atol: Optional[float] = None,
+    equal_nan: bool = False,
+) -> None:
+
+    if base is not None:
+        actual = actual - base
+        expected = expected - base
+
+    torch.testing.assert_close(
+        actual,
+        expected,
+        rtol=rtol,
+        atol=atol,
+        equal_nan=equal_nan,
+        check_dtype=True,
+    )