GAN fixes (#3)

Author: bcebere

decaf/DECAF.py

@@ -1,10 +1,9 @@
 from collections import OrderedDict
-from typing import Any, Optional, Union
+from typing import Any, List, Optional, Union
 
 import networkx as nx
 import numpy as np
 import pytorch_lightning as pl
-import scipy.linalg as slin
 import torch
 import torch.nn as nn
 
@@ -13,13 +12,32 @@
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+def get_nonlin(name: str) -> nn.Module:
+    if name == "none":
+        return nn.Identity()
+    elif name == "elu":
+        return nn.ELU()
+    elif name == "relu":
+        return nn.ReLU()
+    elif name == "leaky_relu":
+        return nn.LeakyReLU()
+    elif name == "selu":
+        return nn.SELU()
+    elif name == "tanh":
+        return nn.Tanh()
+    elif name == "sigmoid":
+        return nn.Sigmoid()
+    elif name == "softmax":
+        return nn.Softmax(dim=-1)
+    else:
+        raise ValueError(f"Unknown nonlinearity {name}")
+
+
 class TraceExpm(torch.autograd.Function):
     @staticmethod
-    def forward(ctx: Any, input: torch.Tensor) -> torch.Tensor:
-        # detach so we can cast to NumPy
-        E = slin.expm(input.detach().numpy())
-        f = np.trace(E)
-        E = torch.from_numpy(E).to(DEVICE)
+    def forward(ctx: Any, data: torch.Tensor) -> torch.Tensor:
+        E = torch.linalg.matrix_exp(data)
+        f = torch.trace(E)
         ctx.save_for_backward(E)
         return torch.as_tensor(f, dtype=input.dtype)
 
@@ -32,49 +50,57 @@ def backward(ctx: Any, grad_output: torch.Tensor) -> torch.Tensor:
 
 trace_expm = TraceExpm.apply
 
-activation_layer = nn.ReLU(inplace=True)
-
 
 class Generator_causal(nn.Module):
     def __init__(
         self,
         z_dim: int,
         x_dim: int,
         h_dim: int,
-        use_mask: bool = False,
         f_scale: float = 0.1,
         dag_seed: list = [],
+        nonlin_out: Optional[List] = None,
     ) -> None:
         super().__init__()
 
+        if nonlin_out is not None:
+            out_dim = 0
+            for act, length in nonlin_out:
+                out_dim += length
+            if out_dim != x_dim:
+                raise RuntimeError("Invalid nonlin_out")
+
         self.x_dim = x_dim
+        self.nonlin_out = nonlin_out
 
         def block(in_feat: int, out_feat: int, normalize: bool = False) -> list:
             layers = [nn.Linear(in_feat, out_feat)]
             if normalize:
                 layers.append(nn.BatchNorm1d(out_feat, 0.8))
-            layers.append(activation_layer)
+            layers.append(nn.ReLU(inplace=True))
             return layers
 
-        self.shared = nn.Sequential(*block(h_dim, h_dim), *block(h_dim, h_dim))
-
-        if use_mask:
+        self.shared = nn.Sequential(*block(h_dim, h_dim), *block(h_dim, h_dim)).to(
+            DEVICE
+        )
 
-            if len(dag_seed) > 0:
-                M_init = torch.rand(x_dim, x_dim) * 0.0
-                M_init[torch.eye(x_dim, dtype=bool)] = 0
-                M_init = torch.rand(x_dim, x_dim) * 0.0
-                for pair in dag_seed:
-                    M_init[pair[0], pair[1]] = 1
+        if len(dag_seed) > 0:
+            M_init = torch.rand(x_dim, x_dim) * 0.0
+            M_init[torch.eye(x_dim, dtype=bool)] = 0
+            M_init = torch.rand(x_dim, x_dim) * 0.0
+            for pair in dag_seed:
+                M_init[pair[0], pair[1]] = 1
 
-                self.M = torch.nn.parameter.Parameter(M_init, requires_grad=False)
-                print("Initialised adjacency matrix as parsed:\n", self.M)
-            else:
-                M_init = torch.rand(x_dim, x_dim) * 0.2
-                M_init[torch.eye(x_dim, dtype=bool)] = 0
-                self.M = torch.nn.parameter.Parameter(M_init)
+            M_init = M_init.to(DEVICE)
+            self.M = torch.nn.parameter.Parameter(M_init, requires_grad=False).to(
+                DEVICE
+            )
         else:
-            self.M = torch.ones(x_dim, x_dim)
+            M_init = torch.rand(x_dim, x_dim) * 0.2
+            M_init[torch.eye(x_dim, dtype=bool)] = 0
+            M_init = M_init.to(DEVICE)
+            self.M = torch.nn.parameter.Parameter(M_init).to(DEVICE)
+
         self.fc_i = nn.ModuleList(
             [nn.Linear(x_dim + 1, h_dim) for i in range(self.x_dim)]
         )
@@ -111,13 +137,28 @@ def sequential(
             x_masked[:, i] = 0.0
             if i in biased_edges:
                 for j in biased_edges[i]:
-                    x_j = x_masked[:, j].detach().numpy()
-                    np.random.shuffle(x_j)
-                    x_masked[:, j] = torch.from_numpy(x_j)
-            out_i = activation_layer(
-                self.fc_i[i](torch.cat([x_masked, z[:, i].unsqueeze(1)], axis=1))
-            )
-            out[:, i] = nn.Sigmoid()(self.fc_f[i](self.shared(out_i))).squeeze()
+                    x_j = x_masked[:, j]
+                    perm = torch.randperm(len(x_j))
+                    x_masked[:, j] = x_j[perm]
+            out_i = self.fc_i[i](torch.cat([x_masked, z[:, i].unsqueeze(1)], axis=1))
+            out_i = nn.ReLU()(out_i)
+            out_i = self.shared(out_i)
+            out_i = self.fc_f[i](out_i).squeeze()
+            out[:, i] = out_i
+
+        if self.nonlin_out is not None:
+            split = 0
+            for act_name, step in self.nonlin_out:
+                activation = get_nonlin(act_name)
+                out[..., split : split + step] = activation(
+                    out[..., split : split + step]
+                )
+
+                split += step
+
+            if split != out.shape[-1]:
+                raise ValueError("Invalid activations")
+
         return out
 
 
@@ -127,9 +168,9 @@ def __init__(self, x_dim: int, h_dim: int) -> None:
 
         self.model = nn.Sequential(
             nn.Linear(x_dim, h_dim),
-            activation_layer,
+            nn.ReLU(),
             nn.Linear(h_dim, h_dim),
-            activation_layer,
+            nn.ReLU(),
             nn.Linear(h_dim, 1),
         )
 
@@ -153,16 +194,14 @@ def __init__(
         batch_size: int = 32,
         lambda_gp: float = 10,
         lambda_privacy: float = 1,
-        d_updates: int = 5,
         eps: float = 1e-8,
         alpha: float = 1,
         rho: float = 1,
         weight_decay: float = 1e-2,
         grad_dag_loss: bool = False,
         l1_g: float = 0,
         l1_W: float = 1,
-        p_gen: float = -1,
-        use_mask: bool = False,
+        nonlin_out: Optional[List] = None,
     ):
         super().__init__()
         self.save_hyperparameters()
@@ -183,8 +222,8 @@ def __init__(
             z_dim=self.z_dim,
             x_dim=self.x_dim,
             h_dim=h_dim,
-            use_mask=use_mask,
             dag_seed=dag_seed,
+            nonlin_out=nonlin_out,
         ).to(DEVICE)
         self.discriminator = Discriminator(x_dim=self.x_dim, h_dim=h_dim).to(DEVICE)
 
@@ -261,21 +300,7 @@ def privacy_loss(
         )
 
     def get_W(self) -> torch.Tensor:
-        if self.hparams.use_mask:
-            return self.generator.M
-        else:
-            W_0 = []
-            for i in range(self.x_dim):
-                weights = self.generator.fc_i[i].weight[
-                    :, :-1
-                ]  # don't take the noise variable's weights
-                W_0.append(
-                    torch.sqrt(
-                        torch.sum((weights) ** 2, axis=0, keepdim=True)
-                        + self.hparams.eps
-                    )
-                )
-            return torch.cat(W_0, axis=0).T
+        return self.generator.M
 
     def dag_loss(self) -> torch.Tensor:
         W = self.get_W()
@@ -288,7 +313,7 @@ def dag_loss(self) -> torch.Tensor:
         )
 
     def sample_z(self, n: int) -> torch.Tensor:
-        return torch.rand(n, self.z_dim) * 2 - 1
+        return torch.randn(n, self.z_dim, device=DEVICE)
 
     @staticmethod
     def l1_reg(model: nn.Module) -> float:
@@ -298,9 +323,10 @@ def l1_reg(model: nn.Module) -> float:
                 l1 = l1 + layer.norm(p=1)
         return l1
 
-    def gen_synthetic(
-        self, x: torch.Tensor, gen_order: Optional[list] = None, biased_edges: dict = {}
-    ) -> torch.Tensor:
+    def gen_synthetic(self, x: torch.Tensor, biased_edges: dict = {}) -> torch.Tensor:
+        self.generator = self.generator.to(DEVICE)
+        x = x.to(DEVICE)
+        gen_order = self.get_gen_order()
         return self.generator.sequential(
             x,
             self.sample_z(x.shape[0]).type_as(x),
@@ -309,15 +335,7 @@ def gen_synthetic(
         )
 
     def get_dag(self) -> np.ndarray:
-        return np.round(self.get_W().detach().numpy(), 3)
-
-    def get_bi_dag(self) -> np.ndarray:
-        dag = np.round(self.get_W().detach().numpy(), 3)
-        bi_dag = np.zeros_like(dag)
-        for i in range(len(dag)):
-            for j in range(i, len(dag)):
-                bi_dag[i][j] = dag[i][j] + dag[j][i]
-        return np.round(bi_dag, 3)
+        return np.round(self.get_W().detach().cpu().numpy(), 3)
 
     def get_gen_order(self) -> list:
         dense_dag = np.array(self.get_dag())
@@ -333,13 +351,8 @@ def training_step(
         # sample noise
         z = self.sample_z(batch.shape[0])
         z = z.type_as(batch)
+        generated_batch = self.generator.sequential(batch, z, self.get_gen_order())
 
-        if self.hparams.p_gen < 0:
-            generated_batch = self.generator.sequential(batch, z, self.get_gen_order())
-        else:  # train simultaneously
-            raise ValueError(
-                "we're not allowing simultaneous generation no more. Set p_gen negative"
-            )
         # train generator
         if optimizer_idx == 0:
             self.iterations_d += 1
@@ -356,12 +369,10 @@ def training_step(
             d_loss += self.hparams.lambda_gp * self.compute_gradient_penalty(
                 batch, generated_batch
             )
+            if torch.isnan(d_loss).sum() != 0:
+                raise ValueError("NaN in the discr loss")
 
-            tqdm_dict = {"d_loss": d_loss.detach()}
-            output = OrderedDict(
-                {"loss": d_loss, "progress_bar": tqdm_dict, "log": tqdm_dict}
-            )
-            return output
+            return d_loss
         elif optimizer_idx == 1:
             # sanity check: keep track of G updates
             self.iterations_g += 1
@@ -382,14 +393,10 @@ def training_step(
             if len(self.dag_seed) == 0:
                 if self.hparams.grad_dag_loss:
                     g_loss += self.gradient_dag_loss(batch, z)
+            if torch.isnan(g_loss).sum() != 0:
+                raise ValueError("NaN in the gen loss")
 
-            tqdm_dict = {"g_loss": g_loss.detach()}
-
-            output = OrderedDict(
-                {"loss": g_loss, "progress_bar": tqdm_dict, "log": tqdm_dict}
-            )
-
-            return output
+            return g_loss
         else:
             raise ValueError("should not get here")
 
@@ -411,7 +418,4 @@ def configure_optimizers(self) -> tuple:
             betas=(b1, b2),
             weight_decay=weight_decay,
         )
-        return (
-            {"optimizer": opt_d, "frequency": self.hparams.d_updates},
-            {"optimizer": opt_g, "frequency": 1},
-        )
+        return [opt_d, opt_g], []

decaf/data.py

@@ -8,11 +8,13 @@
 
 import decaf.logger as log
 
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 class Dataset(torch.utils.data.Dataset):
     def __init__(self, data: list) -> None:
         data = np.array(data, dtype="float32")
-        self.x = torch.from_numpy(data)
+        self.x = torch.from_numpy(data).to(DEVICE)
         self.n_samples = self.x.shape[0]
         log.info("***** DATA ****")
         log.info(f"n_samples = {self.n_samples}")
@@ -49,10 +51,14 @@ def train_dataloader(self) -> DataLoader:
 
     def val_dataloader(self) -> DataLoader:
         return DataLoader(
-            self.data_val, batch_size=self.batch_size, num_workers=self.num_workers
+            self.data_val,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
         )
 
     def test_dataloader(self) -> DataLoader:
         return DataLoader(
-            self.data_test, batch_size=self.batch_size, num_workers=self.num_workers
+            self.data_test,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
         )

decaf/logger.py

@@ -8,7 +8,6 @@
 
 LOG_FORMAT = "[{time}][{process.id}][{level}] {message}"
 
-logger.remove()
 DEFAULT_SINK = "decaf_{time}.log"
 
 

decaf/version.py

@@ -1 +1 @@
-__version__ = "0.1.2"
+__version__ = "0.1.3"

tests/test_decaf.py

@@ -72,11 +72,5 @@ def test_sanity_generate() -> None:
 
     trainer.fit(model, dummy_dm)
 
-    synth_data = (
-        model.gen_synthetic(
-            raw_data, gen_order=model.get_gen_order(), biased_edges=bias_dict
-        )
-        .detach()
-        .numpy()
-    )
+    synth_data = model.gen_synthetic(raw_data, biased_edges=bias_dict)
     assert synth_data.shape[0] == 10