Update examples to work on GPU (where applicable)

Author: alexgdebeer

docs/__init__.py

@@ -79,11 +79,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "torch.set_default_device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "torch.manual_seed(1)\n",
     "set_plot_style()"
    ]

docs/examples/heat/example_inference.ipynb

@@ -13,18 +13,18 @@
 Data = namedtuple("Data", ["X", "y"])
 
 
-def read_credit_data(fname: Path) -> Data:
+def load_credit_data(device: torch.device = torch.device("cpu")) -> Data:
     """Reads in the German credit dataset, then shifts and scales the 
     predictors such that each has a mean of zero and standard deviation 
     of 1, and scales the response variable such that it takes values in 
     {0, 1}.
     """
 
-    with open(fname, "r") as f:
+    with open(DATA_PATH, "r") as f:
         data = [[float(l) for l in line.strip().split()] 
                 for line in f.readlines()]
 
-    data = torch.tensor(data)
+    data = torch.tensor(data, device=device)
     X, y = data[:, :-1], data[:, -1]
 
     mean_X = torch.mean(X, dim=0)
@@ -33,7 +33,4 @@ def read_credit_data(fname: Path) -> Data:
     X = (X - mean_X) / std_X
     y -= 1.0
 
-    return Data(X, y)
-
-
-credit_data = read_credit_data(DATA_PATH)
+    return Data(X, y)

docs/examples/logistic/__init__.py

@@ -42,10 +42,14 @@ def pairplot(
     if labels is None:
         labels = [f"$x_{i+1}$" for i in range(dim)]
 
+    xs = xs.to("cpu")
     if ys is not None:
+        ys = ys.to("cpu")
         xys = torch.vstack((xs, ys))
     else:
         xys = xs.clone()
+    if truth is not None:
+        truth.to("cpu")
 
     if bounds is None:
         samples_min = xys.min(dim=0).values

docs/examples/logistic/example_inference.ipynb

@@ -77,12 +77,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "torch.manual_seed(0)\n",
-    "torch.set_default_device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "set_plot_style()"
    ]
   },
@@ -95,7 +94,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -111,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {

docs/examples/plotting.py

@@ -1,28 +1,23 @@
+from collections import namedtuple
+from pathlib import Path
+
 import torch
 
-from .preconditioner import GammaNormalMapping
-
-
-# Define failure distances (km)
-failure_dists = torch.tensor([
-    6700,  6950,  7820,  8790,  9120,
-    9660,  9820,  11310, 11690, 11850, 
-    11880, 12140, 12200, 12870, 13150, 
-    13330, 13470, 14040, 14300, 17520,
-    17540, 17890, 18420, 18960, 18980,
-    19410, 20100, 20100, 20150, 20320, 
-    20900, 22700, 23490, 26510, 27410, 
-    27490, 27890, 28100
-])
-
-# Define whether or not each observation is right-censored
-censored = torch.tensor([
-    False, True,  True,  True,  False, 
-    True,  True,  True,  True,  True, 
-    True,  True,  False, True,  False, 
-    True,  True,  True,  False, False, 
-    True,  True,  True,  True,  True, 
-    True,  False, True,  True,  True, 
-    False, False, True,  False, True, 
-    False, True,  True
-])
+from .preconditioner import GammaNormalMapping 
+
+
+DATA_PATH = Path(__file__).resolve().parent.joinpath("data")
+FAILURE_DISTS_PATH = DATA_PATH.joinpath("failure_dists.pt")
+CENSORED_PATH = DATA_PATH.joinpath("censored.pt")
+
+
+Data = namedtuple("Data", ["failure_dists", "censored"])
+
+
+def load_shock_data(device: torch.device = torch.device("cpu")) -> Data:
+    """Reads in the data (failure distances (km) and censorship 
+    information) used in @Dolgov2020.
+    """
+    failure_dists = torch.load(FAILURE_DISTS_PATH).to(device=device)
+    censored = torch.load(CENSORED_PATH).to(device=device)
+    return Data(failure_dists, censored)

docs/examples/poisson/example_amortised.ipynb

@@ -3,7 +3,7 @@
 
 import deep_tensor as dt
 
-from examples.shock.prior import GammaDist, GaussianDist
+from .prior import GammaDist, GaussianDist
 
 EPS = torch.finfo(torch.get_default_dtype()).eps
 
@@ -14,8 +14,8 @@ def __init__(
         self, 
         reference: dt.Reference,
         bounds: Tensor, 
-        alpha: float, 
-        gamma: float, 
+        alpha: Tensor, 
+        gamma: Tensor, 
         ms: Tensor, 
         sds: Tensor, 
         dim: int

docs/examples/shock/__init__.py

@@ -87,9 +87,7 @@ def newton(
 
 
 def converged(fs: Tensor, dls: Tensor) -> bool:
-    error_fs = fs.abs()
-    error_dls = dls.abs()
-    return error_fs.max() < 1e-8
+    return fs.abs().max().item() < 1e-8
 
 
 def gaussian_potential(xs: Tensor, mus: Tensor, sds: Tensor) -> Tensor:
@@ -116,10 +114,10 @@ class GammaDist():
     
     """
 
-    def __init__(self, alpha: float, lamb: float, bounds: Tensor):
+    def __init__(self, alpha: Tensor, lamb: Tensor, bounds: Tensor):
 
-        self.alpha = torch.tensor(alpha) 
-        self.lamb = torch.tensor(lamb)
+        self.alpha = alpha
+        self.lamb = lamb
         self.bounds = bounds
         self.Gamma = Gamma(self.alpha, 1.0 / self.lamb)
 
@@ -142,8 +140,8 @@ def func(xs: Tensor) -> Tuple[Tensor, Tensor]:
             zs -= zs_cdf
             return zs, dzdxs
 
-        l0s = torch.full_like(zs_cdf, self.bounds[0])
-        l1s = torch.full_like(zs_cdf, self.bounds[1])
+        l0s = torch.full_like(zs_cdf, self.bounds[0].item())
+        l1s = torch.full_like(zs_cdf, self.bounds[1].item())
         return newton(func, l0s, l1s)