Improve deepsets module

module can now deal with nan values for nonexisting relative states
in case all relative states are nan, the latent representation is zeroed, which is consitent with an empty sum

src/nets/deepsets.py

@@ -17,10 +17,11 @@ class DeepSetsModule(nn.Module):
         """
         super(DeepSetsModule, self).__init__()
         self.input_dim = input_dim
+        self.latent_dim = latent_dim
         self.output_dim = output_dim
-        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, latent_dim)
-        self.rho = Phi(latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
-        self.pooling = torch.sum  # torch.max # torch.mean
+        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, self.latent_dim)
+        self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
+        self.pooling = torch.sum
 
     @staticmethod
     def from_config(config):
@@ -54,18 +55,22 @@ class DeepSetsModule(nn.Module):
     def forward(self, x):
         """
         Args:
-            x (torch.tensor): (batch_size, dynamic_size, input_dim)
+            x (torch.tensor): ([B, ]max_nv, d)
         Returns:
-            y (torch.tensor): (batch_size, output_dim)
+            y (torch.tensor): ([B, ]output_dim)
         """
-        # use negative dynamic_dim since batch dimensions are inserted at the front
-        dynamic_dim = -2
-        # iterate over dynamic dimension to apply phi to every instance
-        latent = tuple(self.phi(instance) for instance in x.unbind(dynamic_dim) if torch.all(~torch.isnan(instance)))
-        # stack outputs of phi
-        latent = torch.stack(latent, dim=dynamic_dim)
-        # apply pooling function to reduce dynamic dimension
-        latent = self.pooling(latent, dim=dynamic_dim)
+        # mask for selecting only those batches and vehicles where all relative states are not nan
+        # shape (B, max_nv)
+        notnan_mask = torch.all(~torch.isnan(x), dim=-1)
+        # create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in
+        latent = torch.zeros([*x.shape[:-1], self.latent_dim])
+        # evaluate phi for all not NaN entries
+        # x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim)
+        latent[notnan_mask] = self.phi(x[notnan_mask])
+
+        # sum over relative state dimension
+        latent = self.pooling(latent, dim=-2)
+        
         # apply rho network
         y = self.rho(latent)
         return y

tests/nets/test_deepsets.py

@@ -2,6 +2,7 @@ import torch
 import random
 from src.nets import deepsets as ds
 import copy
+import numpy as np
 
 ds_config = {
     "input_dim": 5,
@@ -44,34 +45,69 @@ def test_deepsets():
     m = ds.DeepSetsModule.from_config(ds_config)
 
     input_dim = ds_config["input_dim"]
-    n_dynamic = random.randint(5, 15)
-    x = torch.rand(n_dynamic, input_dim)
+    B = 50
+    max_V = 10
+    batch = []
+    for i in range(B):
+        if i==0: # ensure that there is an example with no vehicles
+            n_dynamic = 0
+        elif i==1: # and one with full vehicles
+            n_dynamic = max_V
+        else:
+            n_dynamic = random.randint(1, max_V)
+        x = torch.rand(n_dynamic, input_dim)
+        n_nan = max_V - n_dynamic
+        x = torch.cat([x, torch.zeros(n_nan, input_dim) * np.nan])
+        assert x.shape == torch.Size([max_V, input_dim])
+        batch.append(x)
+    batch = torch.stack(batch)
+    assert batch.shape == torch.Size([B, max_V, input_dim])
 
-    n_batch = 20
-    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+    y = m(batch)
+    assert y.shape == torch.Size([B, ds_config["output_dim"]])
+    assert torch.isnan(y).sum() == 0
 
-    y = m(x)
-    assert y.shape == torch.Size([n_batch, ds_config["output_dim"]])
-
-    for i in range(n_batch):
-        assert torch.allclose(y[i], y[0])
+    for i in range(B):
+        y = m(batch[i])
+        assert y.shape == torch.Size([ds_config["output_dim"]])
+        assert torch.isnan(y).sum() == 0
 
 def test_deepsets_computation():
-    n_dynamic = random.randint(5,15)
-    n_batch = 7
     input_dim = 5
-    output_dim = 3
-    x = torch.rand(n_dynamic, input_dim)
-    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
-    assert x.shape == torch.Size([n_batch, n_dynamic, input_dim])
+    latent_dim = 8
 
-    phi = torch.nn.Linear(input_dim, output_dim)
+    B = 20
+    max_V = 10
+    batch = []
+    for i in range(B):
+        if i==0: # TODO: Set to i==0
+            n_dynamic = 0
+        else:
+            n_dynamic = random.randint(1, max_V)
+        x = torch.rand(n_dynamic, input_dim)
+        n_nan = max_V - n_dynamic
+        x = torch.cat([x, torch.zeros(n_nan, input_dim) * np.nan])
+        assert x.shape == torch.Size([max_V, input_dim])
+        batch.append(x)
+    batch = torch.stack(batch)
+    assert batch.shape == torch.Size([B, max_V, input_dim])
+    x = batch
 
-    y = torch.stack(tuple(phi(instance) for instance in x.unbind(-2)), dim=-2)
-    assert y.shape == torch.Size([n_batch, n_dynamic, output_dim])
+    ### create phi
+    phi = ds.Phi(input_dim, 1, 10, latent_dim)
 
-    y = y.sum(dim=-2)
-    assert y.shape == torch.Size([n_batch, output_dim])
+    max_nv = x.shape[-2]
+    input_mask = ~torch.isnan(x)
+    batch_dynamic_mask = torch.all(input_mask, dim=-1)
+    assert batch_dynamic_mask.shape == x.shape[:-1]
+    batch_mask = torch.all(batch_dynamic_mask, dim=-1)
+    assert batch_mask.shape == x.shape[:-2]
 
-    for i in range(n_batch):
-        assert torch.allclose(y[i], y[0])
+    batch_dims = x.shape[:-2]
+    latent = torch.zeros([*batch_dims, max_nv, latent_dim])
+    latent[batch_dynamic_mask] = phi(x[batch_dynamic_mask])
+    assert x[batch_dynamic_mask].shape == torch.Size([batch_dynamic_mask.sum(), input_dim])
+    assert phi(x[batch_dynamic_mask]).shape == torch.Size([batch_dynamic_mask.sum(), latent_dim])
+
+    latent = latent.sum(dim=-2)
+    assert latent.shape == torch.Size([B, latent_dim])