Shubh/refactor

.gitignore

@@ -14,6 +14,8 @@ docs/ION*
 .pytest_cache
 *.html
 
+*.ipynb
+src/simple_problems
 notebooks/.ipynb_checkpoints/*
 *.out
 .python-version

src/correction_network/attention.py

@@ -0,0 +1,56 @@
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import torch.nn as nn
+
+"""
+Set Attention Block
+(elements, batch, dim_in) -> (elements, batch, dim_out) [Flip elements and batch if batch_first = True]
+"""
+class SAB(nn.Module):
+    def __init__(self, dim_in, dim_out, num_heads, batch_first=False):
+        super().__init__()
+        self.mab = nn.MultiheadAttention(dim_out, num_heads, batch_first=batch_first)
+        self.fc_q = nn.Linear(dim_in, dim_out)
+        self.fc_k = nn.Linear(dim_in, dim_out)
+        self.fc_v = nn.Linear(dim_in, dim_out)
+
+    def forward(self, X, pad_mask=None):
+        Q = self.fc_q(X)
+        K, V = self.fc_k(X), self.fc_v(X)
+        out, wts = self.mab(Q, K, V, key_padding_mask=pad_mask)
+        return out
+
+"""
+Induced Set Attention Block, more efficient than set-attention
+(time, batch, dim_in) -> (time, batch, dim_out) [Flip time and batch if batch_first = True]
+"""    
+class ISAB(nn.Module):
+    def __init__(self, dim_in, dim_out, num_heads, num_inds, batch_first=False):
+        super().__init__()
+        self.I = nn.Parameter(torch.Tensor(1, num_inds, dim_out))
+        nn.init.xavier_uniform_(self.I)
+        self.mab0 = nn.MultiheadAttention(dim_out, num_heads, kdim=dim_in, vdim=dim_out, batch_first=batch_first)
+        self.mab1 = nn.MultiheadAttention(dim_in, num_heads, kdim=dim_out, vdim=dim_out, batch_first=batch_first)
+
+    def forward(self, X, pad_mask=None):
+        # Fix batch_first case
+        H, _ = self.mab0(self.I.repeat(X.size(0), 1, 1), X, X, key_padding_mask=pad_mask)
+        out, _ = self.mab1(X, H, H, key_padding_mask=pad_mask)
+        return out
+
+
+"""
+Transformer Encoder block
+(time, batch, dim_in) -> (time, batch, dim_out) [Flip time and batch if batch_first = True]
+"""
+
+class TEB(nn.Module):
+    def __init__(self, dim, num_heads, num_layers=1, batch_first=False):
+        super().__init__()
+        base_layer = nn.TransformerEncoderLayer(dim, nhead=num_heads, dim_feedforward=2*dim, dropout=0.0, batch_first=batch_first)
+        self.enc = nn.TransformerEncoder(base_layer, num_layers)
+
+    def forward(self, X, pad_mask=None):
+        out = self.enc(X, src_key_padding_mask=pad_mask)
+        return out

src/correction_network/deepsets.py

@@ -0,0 +1,71 @@
+########################################################################
+# Author(s):    Shubh Gupta
+# Date:         21 September 2021
+# Desc:         Network models for GNSS-based position corrections
+########################################################################
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+########################################################
+# DeepSets (src: https://github.com/yassersouri/pytorch-deep-sets)       
+class InvariantModel(nn.Module):
+    def __init__(self, phi: nn.Module, rho: nn.Module):
+        super().__init__()
+        self.phi = phi
+        self.rho = rho
+
+    def forward(self, x):
+        # compute the representation for each data point
+        x = self.phi.forward(x)
+        # sum up the representations
+        x = torch.sum(x, dim=0, keepdim=False)
+        # compute the output
+        out = self.rho.forward(x)
+        return out
+
+class SmallPhi(nn.Module):
+    def __init__(self, input_size: int, output_size: int = 1, hidden_size: int = 10):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+
+        self.fc1 = nn.Linear(self.input_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, self.output_size)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+class SmallRho(nn.Module):
+    def __init__(self, input_size: int, output_size: int = 1, hidden_size: int = 10):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+
+        self.fc1 = nn.Linear(self.input_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, self.output_size)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+class DeepSetModel(nn.Module):
+    def __init__(self, input_size: int, output_size: int = 1, hidden_size: int = 10):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden_size = hidden_size
+
+        phi = SmallPhi(self.input_size, self.hidden_size)
+        rho = SmallPhi(self.hidden_size, self.output_size)
+        self.net = InvariantModel(phi, rho)
+
+    def forward(self, x, pad_mask=None):
+        out = self.net.forward(x)
+        return out

src/correction_network/pooling.py

@@ -0,0 +1,28 @@
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import torch.nn as nn
+
+"""
+Pooling via Multi-headed Attention Block
+(elements, batch, dim_in) -> (1, batch, dim_out) [Flip elements and batch if batch_first = True]
+"""
+class PMA(nn.Module):
+    def __init__(self, dim, num_heads, num_seeds, batch_first=False):
+        super(PMA, self).__init__()
+        if batch_first:
+            self.S = nn.Parameter(torch.Tensor(1, num_seeds, dim))
+        else:
+            self.S = nn.Parameter(torch.Tensor(num_seeds, 1, dim))
+        nn.init.xavier_uniform_(self.S)
+        self.mab = nn.MultiheadAttention(dim, num_heads, batch_first=batch_first)
+        self.batch_first = batch_first
+
+    def forward(self, X, pad_mask=None):
+        if self.batch_first:
+            Q = self.S.repeat(X.size(0), 1, 1)
+        else:
+            Q = self.S.repeat(1, X.size(1), 1)
+        out, _ = self.mab(Q, X, X, key_padding_mask=pad_mask)
+        return out
+

src/correction_network/set_transformer_reduction.py

@@ -0,0 +1,43 @@
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import torch.nn as nn
+from .attention import *
+from .pooling import *
+
+"""
+Modified Set Transformer to reduce a set of features to a fixed-dim output
+(elements, batch, dim_in) -> (1, batch, dim_out) [Flip elements and batch if batch_first = True]
+"""
+class SetTransformerPointOutput(torch.nn.Module):
+    def __init__(self, dim_input, num_outputs, dim_output, dim_hidden=64, num_heads=4, batch_first=False):
+        super().__init__()
+
+        self.feat_in = nn.Sequential(
+                        nn.Linear(dim_input, dim_hidden),
+                        nn.PReLU()
+                    )
+
+        self.enc = TEB(dim_hidden, num_heads, num_layers=2, batch_first=batch_first)
+
+        self.pool = PMA(dim_hidden, num_heads, num_outputs, batch_first=batch_first)
+
+        self.dec = TEB(dim_hidden, num_heads, num_layers=2, batch_first=batch_first)
+
+        self.feat_out = nn.Sequential(
+                    nn.Linear(dim_hidden*num_outputs, dim_output)
+                    )
+        self.batch_first = batch_first
+
+    def forward(self, x, pad_mask=None):
+        x = self.feat_in(x)
+        x = self.enc(x, pad_mask=pad_mask)
+        x = self.pool(x, pad_mask=pad_mask)
+        x = self.dec(x)
+        if self.batch_first:
+            x = x.reshape(x.shape[0], -1)
+        else:
+            x = x.transpose(1, 0, 2)
+            x = x.reshape(x.shape[0], -1)
+        out = self.feat_out(x)
+        return out