Create particle_transfomer.py

particle_transfomer.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Embed(nn.Module):
+    def __init__(self, input_dim, output_dim, normalize_input=False, event_level=False, activation='gelu'):
+        super().__init__()
+
+        self.input_bn = nn.BatchNorm1d(input_dim) if normalize_input else None
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+        self.fc3 = nn.Linear(output_dim, output_dim)
+        self.event_level = event_level
+        
+
+    def forward(self, x):
+        if self.input_bn is not None:
+            # x: (batch, embed_dim, seq_len)
+            x = self.input_bn(x)
+            if not self.event_level: 
+                x = x.permute(2, 0, 1).contiguous()
+
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc3(x))
+
+        return x
+    
+
+class AttBlock(nn.Module):
+    def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'):
+        super(AttBlock, self).__init__()
+
+        self.layer_norm1 = nn.LayerNorm(embed_dims)
+        self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads)
+        self.layer_norm2 = nn.LayerNorm(embed_dims)
+        self.linear1 = nn.Linear(embed_dims, linear_dims1)
+        self.activation = nn.ReLU() if activation == 'relu' else nn.GELU()
+        self.layer_norm3 = nn.LayerNorm(linear_dims1)
+        self.linear2 = nn.Linear(linear_dims1, linear_dims2)
+
+    def forward(self, x, padding_mask=None):
+        # Layer normalization 1
+        x = self.layer_norm1(x)
+
+        if padding_mask is not None:
+        # Assuming mask is 0 for non-padded and 1 for padded elements,
+        # convert it to a boolean tensor with `True` for padded locations.
+            padding_mask = padding_mask.bool()
+
+        
+        x_att, attention = self.multihead_attention(x, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=True)
+        
+        # Skip connection
+        x = x + x_att # Skip connection
+        # Layer normalization 2
+        x = self.layer_norm2(x)
+        # Linear layer and activation
+        x_linear1 = self.activation(self.linear1(x))
+        # Skip connection for the first linear layer
+        x = x + x_linear1
+        # Layer normalization 3
+        x = self.layer_norm3(x_linear1)
+        # Linear layer with specified output dimensions
+        x_linear2 = self.linear2(x)
+        # Skip connection for the second linear layer
+        x = x + x_linear2
+        return x, attention
+
+class ClassBlock(nn.Module):
+    def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'):
+        super(ClassBlock, self).__init__()
+
+        self.layer_norm1 = nn.LayerNorm(embed_dims)
+        self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads)
+        self.layer_norm2 = nn.LayerNorm(embed_dims)
+        self.linear1 = nn.Linear(embed_dims, linear_dims1)
+        self.activation = nn.ReLU() if activation == 'relu' else nn.GELU()
+        self.layer_norm3 = nn.LayerNorm(linear_dims1)
+        self.linear2 = nn.Linear(linear_dims1, linear_dims2)
+
+    def forward(self, x, class_token, padding_mask=None):
+        # Concatenate the class token to the input sequence along the sequence length dimension  
+        x = torch.cat((class_token, x), dim=0)  # (seq_len+1, batch, embed_dim)
+        # Layer normalization 1
+        x = self.layer_norm1(x)
+
+        # Multihead Attention
+        if padding_mask is not None:
+            # Ensure mask has the correct shape for attention
+            padding_mask = torch.cat((torch.zeros_like(padding_mask[:, :1]), padding_mask), dim=1)
+            padding_mask = padding_mask.bool()
+        
+        
+        x_att, attention = self.multihead_attention(class_token, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=False)
+        # Layer normalization 2
+        x = self.layer_norm2(x_att)
+        x = class_token + x  # Skip connection
+        # Linear layer and activation
+        x_linear1 = self.activation(self.linear1(x))
+        # Layer normalization 3
+        x_linear1 = self.layer_norm3(x_linear1)
+        # Linear layer with specified output dimensions
+        x_linear2 = self.linear2(x_linear1 )
+        # Skip connection for the second linear layer
+        x = x + x_linear2
+        return x, attention
+
+class MLPHead(nn.Module):
+    def __init__(self, input_dim, hidden_dim1, hidden_dim2, output_dim):
+        super(MLPHead, self).__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim1)
+        self.fc2 = nn.Linear(hidden_dim1, hidden_dim2)
+        self.fc3 = nn.Linear(hidden_dim2, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+    
+    
+class AnalysisObjectTransformer(nn.Module):
+    def __init__(self, input_dim_obj, input_dim_event, embed_dims, linear_dims1, linear_dims2, mlp_hidden_1, mlp_hidden_2, num_heads=8):
+        super(AnalysisObjectTransformer, self).__init__()
+
+        self.embed_dims = embed_dims
+       
+        # Embedding layer (assumed to be external)
+        self.embedding_layer = Embed(input_dim_obj, embed_dims)
+        self.embedding_layer_event_level = Embed(input_dim_event, embed_dims, event_level=True)
+
+        # Three blocks of self-attention
+        self.block1 = AttBlock(embed_dims, linear_dims1, linear_dims1, num_heads)
+        self.block2 = AttBlock(linear_dims1, linear_dims1, linear_dims1, num_heads)
+        self.block3 = AttBlock(linear_dims1, linear_dims2, linear_dims2, num_heads)
+        self.block5 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)
+        self.block6 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)
+        self.block7 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)
+
+        # Output linear layer and sigmoid activation
+
+        self.mlp = MLPHead(embed_dims + input_dim_event, mlp_hidden_1, mlp_hidden_2, output_dim=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, event_level, mask=None):
+
+        cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims), requires_grad=True)
+        cls_token = nn.init.trunc_normal_(cls_token, std=.02)
+        # Embedding layer
+
+        x = self.embedding_layer(x)
+        x = x.permute(1, 0, 2)
+        
+        attention_weights = []
+
+        # Three blocks of self-attention
+        x, attention = self.block1(x, padding_mask=mask)
+
+        attention_weights.append(attention)
+        x, attention  = self.block2(x, padding_mask=mask)
+        attention_weights.append(attention)
+        x, attention = self.block3(x, padding_mask=mask)
+        attention_weights.append(attention)
+
+        cls_tokens  = cls_token.expand(1, x.size(1), -1)  # (1, N, C)
+        cls_tokens, attention  = self.block5(x, cls_tokens, padding_mask=mask)
+        cls_tokens, attention  = self.block6(x, cls_tokens, padding_mask=mask)
+        cls_tokens, attention  = self.block7(x, cls_tokens, padding_mask=mask)
+
+        x = torch.cat((cls_tokens.squeeze(0), event_level), dim=-1)
+        x = self.mlp(x)
+        output_probabilities = self.sigmoid(x)
+        return output_probabilities, attention_weights