Add SVDropout layer

.gitignore

@@ -136,3 +136,6 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+# Other files
+main.py

README.md

@@ -16,7 +16,7 @@ ETM was originally published by Adji B. Dieng, Francisco J. R. Ruiz, and David M
 With the tools provided here, you can run ETM on your dataset using simple steps.
 
 # Installation
-You can install the package using ```pip``` by running: ```pip install -U embedded_topic_model```
+You can use this package by cloning this repository. Installation via pip will be updated soon.
 
 # Usage
 To use ETM on your corpus, you must first preprocess the documents into a format understandable by the model.
@@ -59,22 +59,22 @@ embeddings_mapping = embedding.create_word2vec_embedding_from_dataset(documents)
 To create and fit the model using the training data, execute:
 
 ```python
-from embedded_topic_model.models.etm import ETM
-
-# Training an ETM instance
-etm_instance = ETM(
+from embedded_topic_model.core.model import ProdEtm, Model
+from embedded_topic_model.core.trainer import Trainer
+
+# Declare model architecture
+prodetm = ProdEtm(
+    len(vocabulary),
+    num_topics=50,
+    train_embeddings=True
+)
+# Declare a trainer to train/eval model
+topic_model = Trainer(
     vocabulary,
-    embeddings=embeddings_mapping, # You can pass here the path to a word2vec file or
-                                   # a KeyedVectors instance
-    num_topics=8,
-    epochs=300,
-    debug_mode=True,
-    train_embeddings=False, # Optional. If True, ETM will learn word embeddings jointly with
-                            # topic embeddings. By default, is False. If 'embeddings' argument
-                            # is being passed, this argument must not be True
+    prodetm
 )
 
-etm_instance.fit(train_dataset)
+topic_model.fit(train_dataset)
 ```
 
 Also, to obtain the topics, topic coherence or topic diversity of the model, you can do as follows:

embedded_topic_model/core/layers.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class SVDropout2D(nn.Module):
+    """A sparse variational dropout apply to a 2D Tensor.
+    """
+    def __init__(self, n_features, dim=1, threshold=0.5):
+        super().__init__()
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.n_features = n_features
+        self.dim = dim
+        self.threshold = threshold / (1-threshold)
+        self.log_sigma = nn.Parameter(torch.Tensor(n_features))
+        self._init_weights()
+
+    def _init_weights(self):
+        self.log_sigma.data.fill_(-5)        
+
+    def forward(self, x):
+        assert x.dim != 2, \
+            "Must be a 2D Tensor"
+        assert x.shape[self.dim] == self.n_features, \
+            "Mismatch tensor shape"
+        if self.dim == 0: x = x.permute(1, 0)
+
+        if self.training:
+            sigma = torch.exp(self.log_sigma)
+            eps = sigma.data.new(sigma.size()).normal_()
+            a = torch.ones_like(sigma) + sigma * eps
+            return torch.matmul(x, torch.diag(a))
+
+        return torch.matmul(x, torch.diag((torch.exp(self.log_sigma) < self.threshold).float()))
+
+    @property
+    def kl_loss(self):
+        k1, k2, k3 = torch.Tensor([0.63576]).to(self.device), torch.Tensor([1.8732]).to(self.device), torch.Tensor([1.48695]).to(self.device)
+        kl = k1 * torch.sigmoid(k2 + k3 * self.log_sigma) - 0.5 * torch.log1p(torch.exp(-self.log_sigma))
+        kl = - kl.mean()
+        return kl
+
+
+class LinearSVD(nn.Module):
+    """A sparse variational dropout apply to linear layer.
+    """
+    def __init__(self, in_features, out_features, threshold=0.5, bias=True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.threshold = threshold / (1-threshold)
+        self._bias = bias
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        self.W = nn.Parameter(torch.Tensor(out_features, in_features))
+        self.log_sigma = nn.Parameter(torch.Tensor(out_features, in_features))
+        if self._bias: self.bias = nn.Parameter(torch.Tensor(1, out_features))
+
+        self._init_weights()
+
+    def _init_weights(self):
+        if self._bias: self.bias.data.zero_()
+        self.W.data.normal_(0, 0.02)
+        self.log_sigma.data.fill_(-5)        
+
+    def forward(self, x):
+        self.log_alpha = self.log_sigma * 2.0 - 2.0 * torch.log(1e-16 + torch.abs(self.W))
+        self.log_alpha = torch.clamp(self.log_alpha, -10, 10) 
+
+        if self.training:
+            if self._bias: lrt_mean =  F.linear(x, self.W) + self.bias
+            else: lrt_mean =  F.linear(x, self.W)
+            lrt_std = torch.sqrt(F.linear(x * x, torch.exp(self.log_sigma * 2.0)) + 1e-8)
+            eps = lrt_std.data.new(lrt_std.size()).normal_()
+            return lrt_mean + lrt_std * eps
+
+        if self._bias: return F.linear(x, self.W * (torch.exp(self.log_alpha) < self.threshold).float()) + self.bias
+        else: return F.linear(x, self.W * (torch.exp(self.log_alpha) < self.threshold).float()) 
+
+    @property
+    def kl_loss(self):
+        k1, k2, k3 = torch.Tensor([0.63576]).to(self.device), torch.Tensor([1.8732]).to(self.device), torch.Tensor([1.48695]).to(self.device)
+        kl = k1 * torch.sigmoid(k2 + k3 * self.log_alpha) - 0.5 * torch.log1p(torch.exp(-self.log_alpha))
+        kl = - kl.mean()
+        return kl
+