Add DL codes
This commit is contained in:
78
src/datasets.py
Normal file
78
src/datasets.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from torch.utils.data import Dataset
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class singlePhotonDataset(Dataset):
|
||||
def __init__(self, sampleList, labelList, sampleRatio):
|
||||
self.sampleFileList = sampleList
|
||||
self.labelFileList = labelList
|
||||
self.sampleRatio = sampleRatio
|
||||
|
||||
for idx, sampleFile in enumerate(self.sampleFileList):
|
||||
if idx == 0:
|
||||
self.samples = np.load(sampleFile)
|
||||
self.labels = np.load(self.labelFileList[idx])
|
||||
else:
|
||||
self.samples = np.concatenate((self.samples, np.load(sampleFile)), axis=0)
|
||||
self.labels = np.concatenate((self.labels, np.load(self.labelFileList[idx])), axis=0)
|
||||
|
||||
### total number of samples
|
||||
self.length = int(self.samples.shape[0] * self.sampleRatio)
|
||||
print(f"Total number of samples: {self.length}")
|
||||
|
||||
def __getitem__(self, index):
|
||||
sample = self.samples[index]
|
||||
sample = np.expand_dims(sample, axis=0)
|
||||
label = self.labels[index]
|
||||
return torch.tensor(sample, dtype=torch.float32), torch.tensor(label, dtype=torch.float32)
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
class doublePhotonDataset(Dataset):
|
||||
def __init__(self, sampleList, labelList, sampleRatio, datasetName, reuselFactor=1):
|
||||
self.sampleFileList = sampleList
|
||||
self.labelFileList = labelList
|
||||
self.sampleRatio = sampleRatio
|
||||
self.datasetName = datasetName
|
||||
|
||||
for idx, sampleFile in enumerate(self.sampleFileList):
|
||||
if idx == 0:
|
||||
self.samples = np.load(sampleFile)
|
||||
self.labels = np.load(self.labelFileList[idx])
|
||||
else:
|
||||
self.samples = np.concatenate((self.samples, np.load(sampleFile)), axis=0)
|
||||
self.labels = np.concatenate((self.labels, np.load(self.labelFileList[idx])), axis=0)
|
||||
|
||||
### total number of samples
|
||||
self.length = int(self.samples.shape[0] * self.sampleRatio) // 2 * reuselFactor
|
||||
print(f"[{self.datasetName} dataset] \t Total number of samples: {self.length}")
|
||||
|
||||
def __getitem__(self, index):
|
||||
sample = np.zeros((8, 8), dtype=np.float32)
|
||||
idx1 = np.random.randint(0, self.samples.shape[0])
|
||||
idx2 = np.random.randint(0, self.samples.shape[0])
|
||||
photon1 = self.samples[idx1]
|
||||
photon2 = self.samples[idx2]
|
||||
|
||||
### random position for photons in
|
||||
pos_x1 = np.random.randint(1, 4)
|
||||
pos_y1 = np.random.randint(1, 4)
|
||||
sample[pos_y1:pos_y1+5, pos_x1:pos_x1+5] += photon1
|
||||
pos_x2 = np.random.randint(1, 4)
|
||||
pos_y2 = np.random.randint(1, 4)
|
||||
sample[pos_y2:pos_y2+5, pos_x2:pos_x2+5] += photon2
|
||||
sample = sample[1:-1, 1:-1] ### sample size: 6x6
|
||||
sample = torch.tensor(sample, dtype=torch.float32).unsqueeze(0) / 1000. ### to keV
|
||||
|
||||
label1 = self.labels[idx1] + np.array([pos_x1, pos_y1, 0, 0]) - 1
|
||||
label2 = self.labels[idx2] + np.array([pos_x2, pos_y2, 0, 0]) - 1
|
||||
# label = np.concatenate((label1, label2), axis=0)
|
||||
if label1[0] < label2[0]:
|
||||
label = np.concatenate((label1, label2), axis=0)
|
||||
else:
|
||||
label = np.concatenate((label2, label1), axis=0)
|
||||
return sample, torch.tensor(label, dtype=torch.float32)
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
190
src/models.py
Normal file
190
src/models.py
Normal file
@@ -0,0 +1,190 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
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 doublePhotonNet_250909(nn.Module):
|
||||
def __init__(self):
|
||||
super(doublePhotonNet_250909, self).__init__()
|
||||
self.conv1 = nn.Conv2d(1, 3, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv2d(3, 5, kernel_size=3, padding=1)
|
||||
self.conv3 = nn.Conv2d(5, 5, kernel_size=3, padding=1)
|
||||
self.fc1 = nn.Linear(5*6*6, 4)
|
||||
# self.fc2 = nn.Linear(50, 4)
|
||||
# 初始化更稳一些
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
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 = F.relu(self.fc1(x))
|
||||
# x = self.fc2(x)
|
||||
x = self.fc1(x)
|
||||
return x
|
||||
class doublePhotonNet_250910(nn.Module):
|
||||
def __init__(self):
|
||||
super(doublePhotonNet_250910, self).__init__()
|
||||
### x shape: (B, 1, 6, 6)
|
||||
self.conv1 = nn.Conv2d(1, 5, kernel_size=5, padding=2) # (B,5,6,6)
|
||||
self.conv2 = nn.Conv2d(5, 10, kernel_size=5, padding=2) # (B,10,6,6)
|
||||
self.conv3 = nn.Conv2d(10, 20, kernel_size=3, padding=0) # (B,20,4,4)
|
||||
self.fc1 = nn.Linear(20*4*4, 4)
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
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 = F.relu(self.fc1(x))
|
||||
# x = self.fc2(x)
|
||||
x = self.fc1(x) * 6
|
||||
return x
|
||||
|
||||
class doublePhotonNet_251001(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# 保持空间分辨率：使用小卷积核 + 无池化
|
||||
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1) # 6x6 -> 6x6
|
||||
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) # 6x6 -> 6x6
|
||||
self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # 6x6 -> 6x6
|
||||
|
||||
# 全局特征提取（替代全连接层）
|
||||
self.global_avg_pool = nn.AdaptiveAvgPool2d((1,1)) # 64x1x1
|
||||
self.global_max_pool = nn.AdaptiveMaxPool2d((1,1)) # 64x1x1
|
||||
|
||||
# 回归头：输出4个坐标 (x1,y1,x2,y2)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(64, 128),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(128, 4), # 直接输出坐标
|
||||
# nn.Sigmoid() # sigmoid leads to overfitting
|
||||
)
|
||||
# for m in self.modules():
|
||||
# if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
|
||||
# if isinstance(m, nn.Linear):
|
||||
# nn.init.xavier_uniform_(m.weight)
|
||||
# nn.init.zeros_(m.bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.relu(self.conv1(x))
|
||||
x = torch.relu(self.conv2(x))
|
||||
x = torch.relu(self.conv3(x))
|
||||
# x = self.global_avg_pool(x).view(x.size(0), -1)
|
||||
x = self.global_max_pool(x).view(x.size(0), -1)
|
||||
coords = self.fc(x)*6
|
||||
return coords # shape: [B, 4]
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class doublePhotonNet_251001_2(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Backbone: deeper + residual-like blocks
|
||||
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
|
||||
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
|
||||
|
||||
# Spatial Attention Module (轻量但有效)
|
||||
self.spatial_attn = nn.Sequential(
|
||||
nn.Conv2d(128, 1, kernel_size=1),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
# Multi-scale feature fusion (optional but helpful)
|
||||
self.reduce1 = nn.Conv2d(32, 32, kernel_size=1) # from conv1
|
||||
self.reduce2 = nn.Conv2d(64, 32, kernel_size=1) # from conv2
|
||||
self.fuse = nn.Conv2d(32*3, 128, kernel_size=1)
|
||||
|
||||
# Global context with both Max and Avg pooling (better than GAP alone)
|
||||
self.global_max_pool = nn.AdaptiveMaxPool2d((1,1))
|
||||
self.global_avg_pool = nn.AdaptiveAvgPool2d((1,1))
|
||||
|
||||
# Enhanced regression head
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(128 * 2, 256), # concat max + avg
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(128, 4),
|
||||
nn.Sigmoid() # output in [0,1]
|
||||
)
|
||||
|
||||
self._init_weights()
|
||||
|
||||
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):
|
||||
# Feature extraction
|
||||
c1 = F.relu(self.conv1(x)) # [B, 32, 6, 6]
|
||||
c2 = F.relu(self.conv2(c1)) # [B, 64, 6, 6]
|
||||
c3 = F.relu(self.conv3(c2)) # [B,128, 6, 6]
|
||||
|
||||
# Spatial attention: highlight photon peaks
|
||||
attn = self.spatial_attn(c3) # [B, 1, 6, 6]
|
||||
c3 = c3 * attn # reweight features
|
||||
|
||||
# (Optional) Multi-scale fusion — uncomment if needed
|
||||
# r1 = F.interpolate(self.reduce1(c1), size=(6,6), mode='nearest')
|
||||
# r2 = self.reduce2(c2)
|
||||
# fused = torch.cat([r1, r2, c3], dim=1)
|
||||
# c3 = self.fuse(fused)
|
||||
|
||||
# Global context: MaxPool better captures peaks, Avg for context
|
||||
g_max = self.global_max_pool(c3).flatten(1) # [B, 128]
|
||||
g_avg = self.global_avg_pool(c3).flatten(1) # [B, 128]
|
||||
global_feat = torch.cat([g_max, g_avg], dim=1) # [B, 256]
|
||||
|
||||
# Regression
|
||||
coords = self.fc(global_feat) * 6.0 # scale to [0,6)
|
||||
return coords # [B, 4]
|
||||
Reference in New Issue
Block a user