refactor: utils

Author: kozistr

tests/test_optimizers.py

@@ -26,12 +26,13 @@
     Shampoo,
 )
 from tests.utils import (
-    LogisticRegression,
     MultiHeadLogisticRegression,
+    build_environment,
     build_lookahead,
     dummy_closure,
     ids,
     make_dataset,
+    tensor_to_numpy,
 )
 
 OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
@@ -85,25 +86,6 @@
 ]
 
 
-def tensor_to_numpy(x: torch.Tensor) -> np.ndarray:
-    return x.detach().cpu().numpy()
-
-
-def build_environment(use_gpu: bool = False) -> Tuple[Tuple[torch.Tensor, torch.Tensor], nn.Module, nn.Module]:
-    torch.manual_seed(42)
-
-    x_data, y_data = make_dataset()
-    model: nn.Module = LogisticRegression()
-    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
-
-    if use_gpu and torch.cuda.is_available():
-        x_data, y_data = x_data.cuda(), y_data.cuda()
-        model = model.cuda()
-        loss_fn = loss_fn.cuda()
-
-    return (x_data, y_data), model, loss_fn
-
-
 @pytest.mark.parametrize('optimizer_fp32_config', OPTIMIZERS, ids=ids)
 def test_f32_optimizers(optimizer_fp32_config):
     (x_data, y_data), model, loss_fn = build_environment()

tests/utils.py

@@ -76,3 +76,22 @@ def build_lookahead(*parameters, **kwargs):
 
 def ids(v) -> str:
     return f'{v[0].__name__}_{v[1:]}'
+
+
+def build_environment(use_gpu: bool = False) -> Tuple[Tuple[torch.Tensor, torch.Tensor], nn.Module, nn.Module]:
+    torch.manual_seed(42)
+
+    x_data, y_data = make_dataset()
+    model: nn.Module = LogisticRegression()
+    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
+
+    if use_gpu and torch.cuda.is_available():
+        x_data, y_data = x_data.cuda(), y_data.cuda()
+        model = model.cuda()
+        loss_fn = loss_fn.cuda()
+
+    return (x_data, y_data), model, loss_fn
+
+
+def tensor_to_numpy(x: torch.Tensor) -> np.ndarray:
+    return x.detach().cpu().numpy()