diff --git a/src/models.py b/src/models.py index 2274a40..55efff4 100644 --- a/src/models.py +++ b/src/models.py @@ -24,95 +24,6 @@ def get_triple_photon_model_class(version): raise ValueError(f"Model class '{class_name}' not found.") return cls -class singlePhotonNet_250909(nn.Module): - def weight_init(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) - - def __init__(self): - super(singlePhotonNet_250909, self).__init__() - self.conv1 = nn.Conv2d(1, 5, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(5, 10, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(10, 20, kernel_size=3, padding=1) - self.fc = nn.Linear(20*5*5, 2) - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - x = self.fc(x) - return x - -class singlePhotonNet_251020(nn.Module): - ''' - Smaller input size (3x3) - ''' - def weight_init(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) - - def __init__(self): - super(singlePhotonNet_251020, self).__init__() - self.conv1 = nn.Conv2d(1, 5, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(5, 10, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(10, 20, kernel_size=3, padding=1) - self.fc = nn.Linear(20*3*3, 2) - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - x = self.fc(x) - return x - -class singlePhotonNet_251022(nn.Module): - ''' - Smaller input size (3x3) - ''' - def weight_init(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) - - def __init__(self): - super(singlePhotonNet_251022, self).__init__() - self.conv1 = nn.Conv2d(3, 5, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(5, 10, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(10, 20, kernel_size=3) - self.fc = nn.Linear(20, 3) - self.weight_init() - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - x = self.fc(x) - return x - -import torch -import torch.nn as nn -import torch.nn.functional as F - class singlePhotonNet_260511(nn.Module): def __init__(self): super(singlePhotonNet_260511, self).__init__() @@ -147,63 +58,7 @@ class singlePhotonNet_260511(nn.Module): coords = self.fc(flat_feat) # [B, 2] return coords -class doublePhotonNet_260608(nn.Module): ## adapted from 260507, removed padding and dropout - def __init__(self, nSize=6): - super().__init__() - self.nSize = nSize - # Backbone: 3 CNN layers - self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=0) - self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=0) - - # Spatial Attention Module - self.spatial_attn = nn.Sequential( - nn.Conv2d(128, 1, kernel_size=1), - nn.Sigmoid() - ) - - self.fc = nn.Sequential( - nn.Linear(128 * (self.nSize-4) * (self.nSize-4), 512), - nn.ReLU(), - nn.Linear(512, 128), - nn.ReLU(), - nn.Linear(128, 4) - ) - - self._init_weights() - self._init_coords() - - def _init_coords(self): - # Create a coordinate grid; moved from dataset generation to model initialization for lower traffic and more flexibility - x = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize) - y = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize) - x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (nSize,nSize), (nSize,nSize) - self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous().to('cuda') # (1, nSize, nSize) - self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous().to('cuda') # (1, nSize, nSize) - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - x = torch.cat((x, self.x_grid.expand(x.size(0), -1, -1, -1), self.y_grid.expand(x.size(0), -1, -1, -1)), dim=1) # [B, 3, nSize, nSize] - c1 = F.relu(self.conv1(x)) # [B, 32, nSize, nSize] - c2 = F.relu(self.conv2(c1)) # [B, 64, nSize-2, nSize-2] - c3 = F.relu(self.conv3(c2)) # [B, 128, nSize-4, nSize-4] - - attn = self.spatial_attn(c3) # [B, 1, nSize-4, nSize-4] - c3 = c3 * attn # [B, 128, nSize-4, nSize-4] - - flat_feat = c3.view(c3.size(0), -1) # [B, 128 * (nSize-4) * (nSize-4)] - - preditions = self.fc(flat_feat) # [B, 4] - return preditions - -class doublePhotonNet_260610(nn.Module): ## adapted from 260507, removed padding and dropout +class doublePhotonNet_260610(nn.Module): def __init__(self): super().__init__() # Backbone: 3 CNN layers @@ -259,68 +114,7 @@ class doublePhotonNet_260610(nn.Module): ## adapted from 260507, removed padding preditions = self.fc(flat_feat) # [B, 4] return preditions -class triplePhotonNet_260529(nn.Module): ## adapted from doublePhotonNet_260507, add one more conv layer and increase capacity of FC layers, for 3-photon pileup with 9x9 input - def __init__(self): - super().__init__() - # Backbone: deeper for 9x9 input containing 3 photons - self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) ### 9x9 - self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) ## 9x9 - self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) ## 9x9 - self.conv4 = nn.Conv2d(128, 128, kernel_size=3) ## 7x7 - - # Spatial Attention Module - self.spatial_attn = nn.Sequential( - nn.Conv2d(128, 1, kernel_size=1), - nn.Sigmoid() - ) - - self.fc = nn.Sequential( - nn.Linear(128 * 7 * 7, 512), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(512, 128), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(128, 6) - ) - - self._init_weights() - self._init_coords() - - def _init_coords(self): - # Create a coordinate grid; moved from dataset generation to model initialization for lower traffic and more flexibility - nSize = 9 # should match the input size of the model - x = np.linspace(-nSize/2. + 0.5, nSize/2. - 0.5, nSize) - y = np.linspace(-nSize/2. + 0.5, nSize/2. - 0.5, nSize) - x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (nSize,nSize), (nSize,nSize) - self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous().to('cuda') # (1, nSize, nSize) - self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous().to('cuda') # (1, nSize, nSize) - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - x = torch.cat((x, self.x_grid.expand(x.size(0), -1, -1, -1), self.y_grid.expand(x.size(0), -1, -1, -1)), dim=1) # [B, 3, 9, 9] - c1 = F.relu(self.conv1(x)) # [B, 32, 9, 9] - c2 = F.relu(self.conv2(c1)) # [B, 64, 9, 9] - c3 = F.relu(self.conv3(c2)) # [B, 128, 9, 9] - c4 = F.relu(self.conv4(c3)) # [B, 128, 7, 7] - - attn = self.spatial_attn(c4) # [B, 1, 7, 7] - c4 = c4 * attn # [B, 128, 7, 7] - - flat_feat = c4.view(c4.size(0), -1) # [B, 6272] - - coords = self.fc(flat_feat) # [B, 6] - return coords - - -class triplePhotonNet_260611(nn.Module): ## adapted from triplePhotonNet_260529 +class triplePhotonNet_260611(nn.Module): def __init__(self): super().__init__() # Backbone: deeper for 9x9 input containing 3 photons