Clean model zoo
This commit is contained in:
+2
-208
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user