add optional rnn and sequential mode to 1d_reg

src/data/gp_curves.py

@@ -40,6 +40,7 @@ class GPCurvesReader(object):
         sigma_scale=1.0,
         random_kernel_parameters=True,
         testing=False,
+        sequential=False
     ):
         """Creates a regression dataset of functions sampled from a GP.
 
@@ -63,6 +64,7 @@ class GPCurvesReader(object):
         self._sigma_scale = sigma_scale
         self._random_kernel_parameters = random_kernel_parameters
         self._testing = testing
+        self._sequential = sequential
 
     def _gaussian_kernel(self, xdata, l1, sigma_f, sigma_noise=2e-2):
         """Applies the Gaussian kernel to generate curve data.
@@ -112,10 +114,10 @@ class GPCurvesReader(object):
         # If we are testing we want to have more targets and have them evenly
         # distributed in order to plot the function.
         if self._testing:
-            num_target = 400
+            num_target = num_context*2
             num_total_points = num_target
             x_values = (
-                torch.arange(-2, 2, 1.0 / 100).unsqueeze(0).repeat(self._batch_size, 1)
+                torch.linspace(-2, 2, num_target).unsqueeze(0).repeat(self._batch_size, 1)
             )
             x_values = x_values.unsqueeze(-1)
         # During training the number of target points and their x-positions are
@@ -127,6 +129,10 @@ class GPCurvesReader(object):
                 torch.rand((self._batch_size, num_total_points, self._x_size)) * 4 - 2
             )
 
+        # For sequential
+        if self._sequential:
+            x_values = x_values.sort(dim=1)[0]
+
         # Set kernel parameters
         # Either choose a set of random parameters for the mini-batch
         if self._random_kernel_parameters:
@@ -160,7 +166,6 @@ class GPCurvesReader(object):
         y_values = torch.matmul(
             cholesky, torch.randn((self._batch_size, self._y_size, num_total_points, 1))
         )
-
         # [batch_size, num_total_points, y_size]
         y_values = y_values.squeeze(3)
         y_values = y_values.permute(0, 2, 1)
@@ -172,6 +177,8 @@ class GPCurvesReader(object):
 
             # Select the observations
             idx = torch.randperm(num_target)
+            if self._sequential:
+                idx = idx.sort()[0]
             context_x = x_values[:, idx[:num_context]]
             context_y = y_values[:, idx[:num_context]]
 

src/models/model.py

@@ -53,11 +53,24 @@ class LatentModel(nn.Module):
                  batchnorm=False,
                  use_lvar=False,
                  attention_layers=2,
+                 use_rnn=False,
                  **kwargs,
                 ):
 
         super(LatentModel, self).__init__()
 
+        self._use_rnn = use_rnn
+
+        if self._use_rnn:
+            self._lstm = nn.LSTM(
+                input_size=x_dim,
+                hidden_size=hidden_dim,
+                num_layers=attention_layers,
+                dropout=dropout,
+                batch_first=True
+            )
+            x_dim = hidden_dim
+
         self._latent_encoder = LatentEncoder(
             x_dim + y_dim,
             hidden_dim=hidden_dim,
@@ -106,6 +119,15 @@ class LatentModel(nn.Module):
     def forward(self, context_x, context_y, target_x, target_y=None):
         num_targets = target_x.size(1)
 
+        if self._use_rnn:
+            # see https://arxiv.org/abs/1910.09323 where x is substituted with h = RNN(x)
+            # x need to be provided as [B, T, H]
+            x = torch.cat([context_x, target_x], dim=1)
+            # h: [B, T, num_direction * H]
+            h, _ = self._lstm(x)
+            context_x = h[:, :context_x.shape[1], :]
+            target_x = h[:, context_x.shape[1]:, :]
+
         dist_prior, log_var_prior = self._latent_encoder(context_x, context_y)
 
         if target_y is not None: