From ffdff12ccb436276dc0c0a64a3711a27074409cf Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Tue, 20 Jul 2021 18:23:51 +0200 Subject: [PATCH 1/3] Update deepsets to deal with nans (first version) --- src/nets/deepsets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py index edbb4a0..f83d1b1 100644 --- a/src/nets/deepsets.py +++ b/src/nets/deepsets.py @@ -61,7 +61,7 @@ class DeepSetsModule(nn.Module): # 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)) + 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 From 61f06c95c208752dab5eb774c37cf71d47f540f4 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Tue, 20 Jul 2021 18:44:07 +0200 Subject: [PATCH 2/3] 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 | 31 ++++++++------ tests/nets/test_deepsets.py | 80 +++++++++++++++++++++++++++---------- 2 files changed, 76 insertions(+), 35 deletions(-) diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py index f83d1b1..e40978e 100644 --- a/src/nets/deepsets.py +++ b/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 diff --git a/tests/nets/test_deepsets.py b/tests/nets/test_deepsets.py index 2321918..d65b899 100644 --- a/tests/nets/test_deepsets.py +++ b/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]) \ No newline at end of file + 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]) From c939a8e13121953ef0026676a2c1bba5d737ff11 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Tue, 20 Jul 2021 18:45:49 +0200 Subject: [PATCH 3/3] Remove todo --- tests/nets/test_deepsets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/nets/test_deepsets.py b/tests/nets/test_deepsets.py index d65b899..7726f57 100644 --- a/tests/nets/test_deepsets.py +++ b/tests/nets/test_deepsets.py @@ -80,7 +80,7 @@ def test_deepsets_computation(): max_V = 10 batch = [] for i in range(B): - if i==0: # TODO: Set to i==0 + if i==0: n_dynamic = 0 else: n_dynamic = random.randint(1, max_V)