diff --git a/src/models.py b/src/models.py index 9bf2f2a..5f9891a 100644 --- a/src/models.py +++ b/src/models.py @@ -203,6 +203,62 @@ class doublePhotonNet_260608(nn.Module): ## adapted from 260507, removed padding preditions = self.fc(flat_feat) # [B, 4] return preditions +class doublePhotonNet_260610(nn.Module): ## adapted from 260507, removed padding and dropout + def __init__(self): + super().__init__() + # Backbone: 3 CNN layers + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) ### 6x6 + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) ### 6x6 + self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=0) ### 4x4 + + # Spatial Attention Module + self.spatial_attn = nn.Sequential( + nn.Conv2d(64, 1, kernel_size=1), + nn.Sigmoid() + ) + + self.fc = nn.Sequential( + nn.Linear(128 * 4 * 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(-3 + 0.5, 3 - 0.5, 6) + y = np.linspace(-3 + 0.5, 3 - 0.5, 6) + 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, 6, 6] + c1 = F.relu(self.conv1(x)) # [B, 32, 6, 6] + c2 = F.relu(self.conv2(c1)) # [B, 64, 6, 6] + + attn = self.spatial_attn(c2) # [B, 1, 6, 6] + c2 = c2 * attn # [B, 64, 6, 6] + + c3 = F.relu(self.conv3(c2)) # [B, 128, 4, 4] + + flat_feat = c3.view(c3.size(0), -1) # [B, 128 * 4 * 4] + + 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__()