Add files via upload

Author: Shengxiang-Lin

Model Implementation/CNN with 3D convolution.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Define the 3D CNN model
+class CNN3D(nn.Module):
+    def __init__(self):
+        super(CNN3D, self).__init__()
+
+        # Define three 3D convolutional layers with different kernel sizes
+        self.conv3x3x3 = nn.Conv3d(1, 2, kernel_size=3, padding=(0,1,1))  # 3x3x3 kernel (3D)
+        self.conv3x5x5 = nn.Conv3d(1, 2, kernel_size=(3,5,5), padding=(0,2,2))  # 3x5x5 kernel
+        self.conv3x7x7 = nn.Conv3d(1, 2, kernel_size=(3,7,7), padding=(0,3,3))  # 3x7x7 kernel
+
+        # Encoder part (2D convolution, batch normalization, and ReLU layers)
+        self.encoder = nn.Sequential(
+            nn.Conv2d(12, 32, kernel_size=5, padding=2),  # First 2D convolution (input 12 channels)
+            nn.BatchNorm2d(32),                          # Batch normalization for stable training
+            nn.ReLU(inplace=True),                       # ReLU activation
+
+            nn.Conv2d(32, 64, kernel_size=5, padding=2), # Second 2D convolution
+            nn.BatchNorm2d(64),                          # Batch normalization
+            nn.ReLU(inplace=True)                        # ReLU activation
+        )
+
+        # Decoder part (two 2D convolutional layers to reconstruct output)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(64, 16, kernel_size=5, padding=2), # First decoder layer
+            nn.Conv2d(16, 1, kernel_size=5, padding=2),  # Second decoder layer (output is 1 channel)
+            nn.Sigmoid()                                 # Sigmoid activation to normalize output
+        )
+
+    def forward(self, x):  # Input x has shape (batch_size, 4, 56, 56) [channels: emcal, hcal, trkn, trkp]
+        x = x.unsqueeze(1)  # Add an extra dimension to x (batch_size, 1, 4, 56, 56) for 3D convolutions
+
+        # Select subsets of the input data for 3D convolutions
+        x_e_h_n = x[:, :, :3, :, :]  # Select channels: emcal, hcal, trkn (first 3 channels)
+        x_e_h_p = x[:, :, [0, 1, 3], :, :]  # Select channels: emcal, hcal, trkp (channels 0, 1, 3)
+
+        # Apply 3D convolutions on these subsets
+        x2 = self.conv3x3x3(x_e_h_n)  # Apply 3x3x3 convolution on emcal, hcal, trkn
+        x3 = self.conv3x5x5(x_e_h_n)  # Apply 3x5x5 convolution on emcal, hcal, trkn
+        x4 = self.conv3x7x7(x_e_h_n)  # Apply 3x7x7 convolution on emcal, hcal, trkn
+
+        x5 = self.conv3x3x3(x_e_h_p)  # Apply 3x3x3 convolution on emcal, hcal, trkp
+        x6 = self.conv3x5x5(x_e_h_p)  # Apply 3x5x5 convolution on emcal, hcal, trkp
+        x7 = self.conv3x7x7(x_e_h_p)  # Apply 3x7x7 convolution on emcal, hcal, trkp
+
+        # Concatenate all the resulting feature maps (12 channels) along the channel dimension
+        x = torch.cat((x2, x3, x4, x5, x6, x7), dim=1).view(-1, 12, 56, 56)
+
+        # Pass through the encoder
+        x = self.encoder(x)
+
+        # Pass through the decoder to reconstruct the output
+        x = self.decoder(x)
+
+        return x  # Output shape: (batch_size, 1, 56, 56)

Model Implementation/CNN with self-attention.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Define the Self-Attention layer
+class SelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super(SelfAttention, self).__init__()
+        # Convolutional layers for generating query, key, and value matrices from the input feature maps
+        self.query = nn.Conv2d(in_channels, in_channels, kernel_size=1)  # For the query matrix (Q)
+        self.key = nn.Conv2d(in_channels, in_channels, kernel_size=1)    # For the key matrix (K)
+        self.value = nn.Conv2d(in_channels, in_channels, kernel_size=1)  # For the value matrix (V)
+        # Softmax layer to normalize the attention scores
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        batch_size, C, H, W = x.size()  # Get the dimensions of the input tensor
+
+        # Generate query, key, and value matrices by applying respective convolutions
+        queries = self.query(x).view(batch_size, C, -1)  # Reshape to (B, C, H*W)
+        keys = self.key(x).view(batch_size, C, -1)       # Reshape to (B, C, H*W)
+        values = self.value(x).view(batch_size, C, -1)   # Reshape to (B, C, H*W)
+
+        # Compute the attention scores using matrix multiplication between queries and keys
+        attention_scores = torch.bmm(queries.permute(0, 2, 1), keys)  # Output: (B, H*W, H*W)
+        attention_scores = self.softmax(attention_scores)  # Apply softmax to get the attention weights
+
+        # Multiply values with attention scores and reshape back to the original size
+        out = torch.bmm(values, attention_scores.permute(0, 2, 1))  # Output: (B, C, H*W)
+        return out.view(batch_size, C, H, W)  # Reshape to (B, C, H, W) without changing the original shape
+
+# Define the CNN with Self-Attention model
+class CNNattention(nn.Module):
+    def __init__(self, in_channels, out_channels=1):  # Default output channels set to 1
+        super(CNNattention, self).__init__()
+        # Encoder part (downsampling with convolutional layers)
+        self.enc1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)  # First encoding layer, input has 4 channels
+        self.enc2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)          # Second encoding layer
+        self.enc3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)         # Third encoding layer
+
+        # Self-Attention layer applied after the encoder
+        self.attention = SelfAttention(256)  # Self-Attention applied to 256 channels
+
+        # Decoder part (upsampling with convolutional layers)
+        self.dec1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)  # First decoding layer
+        self.dec2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)   # Second decoding layer
+        self.dec3 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)  # Final layer, output has 1 channel
+        nn.Sigmoid()  # Sigmoid activation (not used in forward)
+
+    def forward(self, x):
+        # Encoder: apply convolutional layers and relu activation
+        enc1 = F.relu(self.enc1(x))  # (4, 56, 56) -> (64, 56, 56)
+        enc2 = F.relu(self.enc2(F.max_pool2d(enc1, 2)))  # Downsampling: (64, 28, 28) -> (128, 28, 28)
+        enc3 = F.relu(self.enc3(F.max_pool2d(enc2, 2)))  # Downsampling: (128, 14, 14) -> (256, 14, 14)
+
+        # Self-Attention layer to refine feature maps
+        attn_out = self.attention(enc3)  # Attention output: (256, 14, 14)
+
+        # Decoder: upsample and apply convolutional layers
+        dec1 = F.relu(self.dec1(F.upsample(attn_out, scale_factor=2, mode='bilinear', align_corners=False)))  # Upsample: (128, 28, 28)
+        dec2 = F.relu(self.dec2(F.upsample(dec1, scale_factor=2, mode='bilinear', align_corners=False)))  # Upsample: (64, 56, 56)
+        out = self.dec3(dec2)  # Final output: (64, 56, 56) -> (1, 56, 56)
+
+        return out  # Return the final output

Model Implementation/CNN.py

@@ -0,0 +1,48 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+# Define the CNN model
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+
+        # Three convolutional layers with different kernel sizes (3x3, 5x5, 7x7)
+        self.cov3x3 = nn.Conv2d(4, 2, kernel_size=3, padding=1)  # 3x3 kernel with padding for same size output
+        self.cov5x5 = nn.Conv2d(4, 2, kernel_size=5, padding=2)  # 5x5 kernel with padding for same size output
+        self.cov7x7 = nn.Conv2d(4, 2, kernel_size=7, padding=3)  # 7x7 kernel with padding for same size output
+
+        # Encoder part (a sequence of convolutional, batch normalization, and ReLU layers)
+        self.encoder = nn.Sequential(
+            nn.Conv2d(6, 32, kernel_size=5, padding=2),  # First layer after concatenation of cov layers
+            nn.BatchNorm2d(32),                          # Batch normalization for stable training
+            nn.ReLU(inplace=True),                       # ReLU activation function
+
+            nn.Conv2d(32, 64, kernel_size=5, padding=2), # Second convolutional layer
+            nn.BatchNorm2d(64),                          # Batch normalization
+            nn.ReLU(inplace=True)                        # ReLU activation
+        )
+
+        # Decoder part (two convolutional layers for reconstruction)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(64, 16, kernel_size=5, padding=2), # First layer of decoder
+            nn.Conv2d(16, 1, kernel_size=5, padding=2),  # Second layer, outputting one channel
+            nn.Sigmoid()                                 # Sigmoid activation for the output (for normalization)
+        )
+
+    def forward(self, x):  # Input x is a 4-channel image, i.e., concatenated (emcal, hcal, trkn, trkp)
+        # Apply three convolutional layers with different kernel sizes
+        x1 = self.cov3x3(x)  # Apply 3x3 convolution
+        x2 = self.cov5x5(x)  # Apply 5x5 convolution
+        x3 = self.cov7x7(x)  # Apply 7x7 convolution
+
+        # Concatenate the outputs of the three convolutional layers along the channel dimension
+        x = torch.cat((x1, x2, x3), dim=1)  # (2+2+2 = 6 channels)
+
+        # Pass the concatenated output through the encoder
+        x = self.encoder(x)
+
+        # Pass the encoded features through the decoder to reconstruct the output
+        x = self.decoder(x)
+
+        return x  # Final output (1 channel, 56x56)

Model Implementation/RNN.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+
+class RNN(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers=2):
+        super(RNN, self).__init__()
+        # Initialize the encoder LSTM layer
+        self.encoder = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
+        # Initialize the decoder LSTM layer
+        self.decoder = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
+        # Initialize a fully connected layer to map the hidden states to the output dimension
+        self.fc = nn.Linear(hidden_dim, output_dim)
+        
+    def forward(self, x):
+        # Forward pass through the encoder LSTM
+        # enc_out: Encoded output sequences
+        # hidden: Hidden states of the encoder
+        enc_out, hidden = self.encoder(x)
+        
+        # Forward pass through the decoder LSTM
+        # dec_out: Decoded output sequences
+        # _ : Decoder hidden states (not used here)
+        dec_out, _ = self.decoder(enc_out, hidden)
+        
+        # Apply the fully connected layer to map decoder outputs to the desired output dimension
+        out = self.fc(dec_out)
+        return out

Model Implementation/U-Net with self-attention.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Define a Self-Attention layer
+class SelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super(SelfAttention, self).__init__()
+        # Define the layers for query, key, and value
+        self.query = nn.Conv2d(in_channels, in_channels, kernel_size=1)  # 1x1 convolution to generate queries
+        self.key = nn.Conv2d(in_channels, in_channels, kernel_size=1)    # 1x1 convolution to generate keys
+        self.value = nn.Conv2d(in_channels, in_channels, kernel_size=1)  # 1x1 convolution to generate values
+        self.softmax = nn.Softmax(dim=-1)  # Softmax for attention scores
+
+    def forward(self, x):
+        batch_size, C, H, W = x.size()
+        # Generate queries, keys, and values
+        queries = self.query(x).view(batch_size, C, -1)  # (B, C, H*W)
+        keys = self.key(x).view(batch_size, C, -1)      # (B, C, H*W)
+        values = self.value(x).view(batch_size, C, -1)  # (B, C, H*W)
+
+        # Compute attention scores
+        attention_scores = torch.bmm(queries.permute(0, 2, 1), keys)  # (B, H*W, H*W)
+        attention_scores = self.softmax(attention_scores)  # Apply softmax
+
+        # Compute the attention output
+        out = torch.bmm(values, attention_scores.permute(0, 2, 1))  # (B, C, H*W)
+        return out.view(batch_size, C, H, W)  # Reshape to original dimensions
+
+# Define the U-Net with Transformer integration
+class UNetTransformer(nn.Module):
+    def __init__(self, in_channels, out_channels=1):  # Output channels set to 1
+        super(UNetTransformer, self).__init__()
+        # Encoder part
+        self.enc1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)  # Input has 4 channels
+        self.enc2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.enc3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        # Self-attention layer
+        self.attention = SelfAttention(256)  # Input has 256 channels
+        # Decoder part
+        self.dec1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.conv1 = nn.Conv2d(256, 128, kernel_size=1)  # Adjust number of channels after concatenation
+        self.dec2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.dec3 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)  # Final output layer
+        self.conv2 = nn.Conv2d(128, 64, kernel_size=1)  # Adjust number of channels after concatenation
+
+    def forward(self, x):
+        # Encoding
+        enc1 = F.relu(self.enc1(x))  # 4x56x56 -> 64x56x56
+        enc2 = F.relu(self.enc2(F.max_pool2d(enc1, 2)))  # 64x28x28 -> 128x28x28
+        enc3 = F.relu(self.enc3(F.max_pool2d(enc2, 2)))  # 128x14x14 -> 256x14x14
+
+        # Self-attention
+        attn_out = self.attention(enc3)  # 256x14x14 -> 256x14x14
+
+        # Decoding
+        dec1 = F.relu(self.dec1(F.upsample(attn_out, scale_factor=2, mode='bilinear', align_corners=False)))  # 256x14x14 -> 128x28x28
+        dec1 = torch.cat([dec1, enc2], dim=1)  # Concatenate with corresponding encoder output
+        dec1 = self.conv1(dec1)  # Adjust channels after concatenation
+        dec2 = F.relu(self.dec2(F.upsample(dec1, scale_factor=2, mode='bilinear', align_corners=False)))  # 128x28x28 -> 64x56x56
+        dec2 = torch.cat([dec2, enc1], dim=1)  # Concatenate with corresponding encoder output
+        dec2 = self.conv2(dec2)  # Adjust channels after concatenation
+        out = self.dec3(dec2)  # 64x56x56 -> 1x56x56
+        
+        return out  # Output the final result

Model Implementation/U-Net.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Basic convolutional block with two convolutional layers, batch normalization, dropout, and LeakyReLU activation
+class Conv(nn.Module):
+    def __init__(self, C_in, C_out):
+        super(Conv, self).__init__()
+        self.layer = nn.Sequential(
+            # First 3x3 convolution layer with batch normalization, dropout, and activation
+            nn.Conv2d(C_in, C_out, kernel_size=3, stride=1, padding=1),  # 3x3 convolution, padding=1 keeps size unchanged
+            nn.BatchNorm2d(C_out),
+            nn.Dropout(0.3),
+            nn.LeakyReLU(inplace=True),
+            # Second 3x3 convolution layer with batch normalization, dropout, and activation
+            nn.Conv2d(C_out, C_out, kernel_size=3, stride=1, padding=1),  # Again, keep size unchanged
+            nn.BatchNorm2d(C_out),
+            nn.Dropout(0.4),
+            nn.LeakyReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+# Downsampling block to reduce the feature map size by half using a 2x2 convolution with stride 2
+class DownSampling(nn.Module):
+    def __init__(self, C_in, C_out):
+        super(DownSampling, self).__init__()
+        self.Down = nn.Sequential(
+            # 2x2 convolution with stride 2 reduces the spatial dimensions
+            nn.Conv2d(C_in, C_out, kernel_size=2, stride=2),  # 2x2 convolution, stride=2 reduces size by half
+            nn.LeakyReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.Down(x)
+
+# Upsampling block to increase the feature map size by a factor of 2 using nearest neighbor interpolation
+class UpSampling(nn.Module):
+    def __init__(self, C_in, C_out):
+        super(UpSampling, self).__init__()
+        # 1x1 convolution to adjust the number of channels after upsampling
+        self.Up = nn.Conv2d(C_in, C_out, kernel_size=1)  # 1x1 convolution to change the number of channels
+
+    def forward(self, x, r):
+        # Upsample using nearest neighbor interpolation
+        up = F.interpolate(x, scale_factor=2, mode='nearest')  # Use nearest neighbor interpolation to upsample
+        x = self.Up(up)  # Change the number of output channels
+        x = torch.cat([x, r], dim=1)  # Concatenate the upsampled feature map with the skip connection
+        return x
+
+# U-Net architecture for image segmentation or similar tasks
+class UNet(nn.Module):
+    def __init__(self):
+        super(UNet, self).__init__()
+        # Initial convolution block with 4 input channels (for example, emcal, hcal, trkn, trkp)
+        self.C1 = Conv(4, 64)  # First convolution block, input has 4 channels
+        # First downsampling layer to reduce spatial size
+        self.D1 = DownSampling(64, 128)  # First downsampling layer
+        # Second convolution block
+        self.C2 = Conv(128, 128)  # Second convolution block
+        # Skipping a second downsampling and upsampling for simplicity in this version
+        # Final upsampling layer to increase the size back and perform concatenation with the previous layer
+        self.U2 = UpSampling(128, 64)  # Second upsampling layer
+        # Final convolution block after concatenation
+        self.C5 = Conv(128, 64)  # Final convolution block
+        # Prediction layer that outputs a single channel
+        self.pred = nn.Conv2d(64, 1, kernel_size=3, padding=1)  # Output layer, predicts 1 channel
+        # Apply sigmoid activation for final output
+        self.sigmoid = nn.Sigmoid()  # Sigmoid activation for final output
+
+    def forward(self, x):
+        # Encoder part: First convolution and downsampling
+        R1 = self.C1(x)  # First convolution block
+        R2 = self.C2(self.D1(R1))  # Downsampling followed by the second convolution block
+        
+        # Decoder part: Upsample and concatenate with previous layer (skip connection)
+        up2 = self.U2(R2, R1)  # Second upsampling with skip connection to the first convolution block
+        
+        # Final convolution and output
+        c = self.C5(up2)  # Final convolution block
+        return self.sigmoid(self.pred(c))  # Output prediction with sigmoid activation