New model and loss_fn to improve uniformity
Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
+52
-1
@@ -326,4 +326,55 @@ class doublePhotonNet_251124(nn.Module): ### adapted for 7x7 input from 251001_2
|
||||
|
||||
# Regression
|
||||
coords = self.fc(global_feat)
|
||||
return coords # [B, 4]
|
||||
return coords # [B, 4]
|
||||
|
||||
class doublePhotonNet_260507(nn.Module): ## adapted from 251124, removed max pooling, added more capacity in FC layers
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Backbone 保持不变
|
||||
self.conv1 = nn.Conv2d(3, 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()
|
||||
)
|
||||
|
||||
# 我们移除了全局池化层 (Global Pooling)。
|
||||
# 7x7 的空间特征图有 128 个通道,展平后是 128 * 7 * 7 = 6272 个特征。
|
||||
# 这使得全连接层能够直接“看到”精确的空间分布规律。
|
||||
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, 4)
|
||||
)
|
||||
|
||||
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):
|
||||
c1 = F.relu(self.conv1(x)) # [B, 32, 7, 7]
|
||||
c2 = F.relu(self.conv2(c1)) # [B, 64, 7, 7]
|
||||
c3 = F.relu(self.conv3(c2)) # [B, 128, 7, 7]
|
||||
|
||||
attn = self.spatial_attn(c3) # [B, 1, 7, 7]
|
||||
c3 = c3 * attn # [B, 128, 7, 7]
|
||||
|
||||
# 直接展平 (Flatten) 而不是池化
|
||||
flat_feat = c3.view(c3.size(0), -1) # [B, 6272]
|
||||
|
||||
coords = self.fc(flat_feat) # [B, 4]
|
||||
return coords
|
||||
Reference in New Issue
Block a user