AncTe/ginka-generator
1
0
Fork 0
mirror of https://github.com/unanmed/ginka-generator.git synced 2026-05-18 15:41:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 改进判别器与生成器网络

Browse Source
This commit is contained in:
unanmed 2025-03-25 21:12:30 +08:00
parent 87f48dc8ee
commit 8872de3f13
11 changed files with 116 additions and 135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
data/src/topology/similarity.ts
View File

@ -59,9 +59,10 @@ function weisfeilerLehmanIteration(
const neighborLabels = node.neighbors
.map(n => n.currentLabel)
.sort();
const compositeLabel = `${node.currentLabel}|${neighborLabels.join(
','
)}`;
)}`.slice(0, 4096);
newLabels.push(compositeLabel);
});

4
ginka/model/loss.py
View File

@ -294,7 +294,7 @@ def entrance_spatial_constraint(
return total_loss
class GinkaLoss(nn.Module):
def __init__(self, minamo: MinamoModel, weight=[0.5, 0.15, 0.15, 0.1, 0.1]):
def __init__(self, minamo: MinamoModel, weight=[0.5, 0.1, 0.1, 0.2, 0.1]):
"""Ginka Model 损失函数部分
Args:
@ -335,7 +335,7 @@ class GinkaLoss(nn.Module):
losses = [
minamo_loss * self.weight[0],
border_loss * self.weight[1] * 0.1,
border_loss * self.weight[1],
entrance_loss * self.weight[2],
count_loss * self.weight[3],
illegal_loss * self.weight[4]

19
ginka/model/model.py
View File

@ -16,15 +16,18 @@ class GinkaModel(nn.Module):
nn.BatchNorm1d(fc_dim),
nn.ReLU()
)
self.deconv_layers = nn.Sequential(
nn.ConvTranspose2d(base_ch*8, base_ch*4, kernel_size=4, stride=2, padding=1), # Upsample 2x
nn.BatchNorm2d(base_ch*4),
self.upsample = nn.Sequential(
nn.Conv2d(base_ch*8, base_ch*16, kernel_size=3, stride=1, padding=1),
nn.PixelShuffle(2),
nn.InstanceNorm2d(base_ch*4),
nn.ReLU(),
nn.ConvTranspose2d(base_ch*4, base_ch*2, kernel_size=4, stride=2, padding=1), # Upsample 2x
nn.BatchNorm2d(base_ch*2),
nn.Conv2d(base_ch*4, base_ch*8, kernel_size=3, stride=1, padding=1),
nn.PixelShuffle(2),
nn.InstanceNorm2d(base_ch*2),
nn.ReLU(),
nn.ConvTranspose2d(base_ch*2, base_ch, kernel_size=4, stride=2, padding=1), # Upsample 2x
nn.BatchNorm2d(base_ch),
nn.Conv2d(base_ch*2, base_ch*4, kernel_size=3, stride=1, padding=1),
nn.PixelShuffle(2),
nn.InstanceNorm2d(base_ch),
nn.ReLU(),
)
self.unet = GinkaUNet(base_ch, num_classes)
@ -40,7 +43,7 @@ class GinkaModel(nn.Module):
"""
x = self.fc(feat)
x = x.view(-1, self.base_ch*8, 4, 4)
x = self.deconv_layers(x)
x = self.upsample(x)
x = self.unet(x)
x = F.interpolate(x, (13, 13), mode='bilinear')
return x, F.softmax(x, dim=1)

59
ginka/model/unet.py
View File

@ -1,39 +1,50 @@
import torch
import torch.nn as nn
from shared.attention import CBAM
import torch.nn.functional as F
from shared.attention import CBAM, SEBlock
class GinkaEncoder(nn.Module):
"""编码器（下采样）部分"""
def __init__(self, in_channels, out_channels):
def __init__(self, in_channels, out_channels, attention=False, block='CBAM'):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.InstanceNorm2d(out_channels),
nn.GELU(),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
# CBAM(out_channels),
nn.GELU()
nn.InstanceNorm2d(out_channels),
)
# 注意力
if attention:
if block == 'CBAM':
self.conv.append(CBAM(out_channels))
elif block == 'SEBlock':
self.conv.append(SEBlock(out_channels))
self.conv.append(nn.GELU())
self.pool = nn.MaxPool2d(2)
def forward(self, x):
x_res = self.conv(x)
x_down = self.pool(x_res)
return x_down, x_res
x_res = self.conv(x)
x_down = self.pool(x_res)
return x_down, x_res
class GinkaDecoder(nn.Module):
"""解码器（上采样）部分"""
def __init__(self, in_channels, out_channels):
def __init__(self, in_channels, out_channels, attention=False, block='CBAM'):
super().__init__()
self.upsample = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
self.conv = nn.Sequential(
nn.Conv2d(in_channels + out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
# CBAM(out_channels),
nn.GELU()
nn.InstanceNorm2d(out_channels),
)
# 注意力
if attention:
if block == 'CBAM':
self.conv.append(CBAM(out_channels))
elif block == 'SEBlock':
self.conv.append(SEBlock(out_channels))
self.conv.append(nn.GELU())
def forward(self, x, skip):
x = self.upsample(x)
@ -46,10 +57,11 @@ class GinkaBottleneck(nn.Module):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.InstanceNorm2d(out_channels),
nn.GELU(),
SEBlock(out_channels),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.InstanceNorm2d(out_channels),
nn.GELU(),
)
@ -61,21 +73,20 @@ class GinkaUNet(nn.Module):
"""Ginka Model UNet 部分
"""
super().__init__()
self.down1 = GinkaEncoder(in_ch, in_ch*2)
self.down2 = GinkaEncoder(in_ch*2, in_ch*4)
self.down3 = GinkaEncoder(in_ch*4, in_ch*8)
self.down4 = GinkaEncoder(in_ch*8, in_ch*16)
self.down1 = GinkaEncoder(in_ch, in_ch*2, attention=True)
self.down2 = GinkaEncoder(in_ch*2, in_ch*4, attention=True)
self.down3 = GinkaEncoder(in_ch*4, in_ch*8, attention=True, block='SEBlock')
self.down4 = GinkaEncoder(in_ch*8, in_ch*16, attention=True, block='SEBlock')
self.bottleneck = GinkaBottleneck(in_ch*16, in_ch*16)
self.up1 = GinkaDecoder(in_ch*16, in_ch*8)
self.up2 = GinkaDecoder(in_ch*8, in_ch*4)
self.up3 = GinkaDecoder(in_ch*4, in_ch*2)
self.up4 = GinkaDecoder(in_ch*2, in_ch)
self.up1 = GinkaDecoder(in_ch*16, in_ch*8, attention=True, block='SEBlock')
self.up2 = GinkaDecoder(in_ch*8, in_ch*4, attention=True, block='SEBlock')
self.up3 = GinkaDecoder(in_ch*4, in_ch*2, attention=True)
self.up4 = GinkaDecoder(in_ch*2, in_ch, attention=True)
self.final = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 1),
# nn.Softmax(dim=1) # 适用于分类任务
)
def forward(self, x):

1
ginka/train.py
View File

@ -95,7 +95,6 @@ def train():
# for name, param in model.named_parameters():
# if param.grad is not None:
# print(f"{name}: grad_mean={param.grad.abs().mean():.3e}, max={param.grad.abs().max():.3e}")
avg_loss = total_loss / len(dataloader)
tqdm.write(f"[INFO {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Epoch: {epoch + 1} | loss: {avg_loss:.6f} | lr: {(optimizer.param_groups[0]['lr']):.6f}")

8
minamo/model/loss.py
View File

@ -1,16 +1,16 @@
import torch.nn as nn
class MinamoLoss(nn.Module):
def __init__(self, vision_weight=0.4, topo_weight=0.6):
def __init__(self, vision_weight=1, topo_weight=0):
super().__init__()
self.vision_weight = vision_weight
self.topo_weight = topo_weight
self.mse = nn.MSELoss()
self.loss = nn.L1Loss()
def forward(self, vis_pred, topo_pred, vis_true, topo_true):
# print(vis_pred.shape, topo_pred.shape, vis_true.shape, topo_true.shape)
# print(vis_pred[0].item(), topo_pred[0].item(), vis_true[0].item(), topo_true[0].item())
vis_loss = self.mse(vis_pred, vis_true)
topo_loss = self.mse(topo_pred, topo_true)
vis_loss = self.loss(vis_pred, vis_true)
topo_loss = self.loss(topo_pred, topo_true)
# print(vis_loss.item(), topo_loss.item())
return self.vision_weight * vis_loss + self.topo_weight * topo_loss

42
minamo/model/topo.py
View File

@ -1,8 +1,7 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from torch_geometric.nn import global_mean_pool, TopKPooling, GATConv
from torch_geometric.nn import GATConv, AttentionalAggregation, global_max_pool
from torch_geometric.data import Data
class MinamoTopoModel(nn.Module):
@ -11,23 +10,19 @@ class MinamoTopoModel(nn.Module):
):
super().__init__()
# 传入 softmax 概率值，直接映射
self.input_proj = torch.nn.Linear(tile_types, emb_dim)
self.input_proj = nn.Linear(tile_types, emb_dim)
# 图卷积层
self.conv1 = GATConv(emb_dim, hidden_dim*2, heads=8, dropout=0.2)
self.conv2 = GATConv(hidden_dim*16, hidden_dim*4, heads=4)
self.conv_ins2 = GATConv(hidden_dim*16, hidden_dim*4, heads=4)
self.conv_ins1 = GATConv(hidden_dim*16, hidden_dim*8, heads=2)
self.conv3 = GATConv(hidden_dim*16, out_dim, concat=False)
self.conv2 = GATConv(hidden_dim*16, hidden_dim*2, heads=8)
self.conv3 = GATConv(hidden_dim*16, hidden_dim*2, heads=8)
self.conv4 = GATConv(hidden_dim*16, out_dim, heads=1)
# 正则化
self.norm1 = nn.LayerNorm(hidden_dim*16)
self.norm2 = nn.LayerNorm(hidden_dim*16)
self.norm_ins2 = nn.LayerNorm(hidden_dim*16)
self.norm_ins1 = nn.LayerNorm(hidden_dim*16)
self.norm3 = nn.LayerNorm(out_dim)
self.norm3 = nn.LayerNorm(hidden_dim*16)
self.norm4 = nn.LayerNorm(out_dim)
# 池化层
self.pool = TopKPooling(out_dim, ratio=0.8) # 保留80%关键节点
self.drop = nn.Dropout(0.3)
# 增强MLP
@ -37,30 +32,25 @@ class MinamoTopoModel(nn.Module):
def forward(self, graph: Data):
x = self.input_proj(graph.x)
# identity = x
x = self.conv1(x, graph.edge_index)
x = F.elu(self.norm1(x))
x = F.relu(self.norm1(x))
x = self.conv2(x, graph.edge_index)
x = F.elu(self.norm2(x))
x = self.conv_ins2(x, graph.edge_index)
x = F.elu(self.norm_ins2(x))
x = self.conv_ins1(x, graph.edge_index)
x = F.elu(self.norm_ins1(x))
x = F.relu(self.norm2(x))
x = self.conv3(x, graph.edge_index)
x = F.elu(self.norm3(x))
x = F.relu(self.norm3(x))
# 分层池化
x = self.conv4(x, graph.edge_index)
x = F.relu(self.norm4(x))
# 池化
x = self.drop(x)
# x, _, _, batch, _, _ = self.pool(x, graph.edge_index, batch=graph.batch)
x = global_mean_pool(x, graph.batch)
x = global_max_pool(x, graph.batch)
topo_vec = self.fc(x)
# 增强MLP
# 归一化
return F.normalize(topo_vec, p=2, dim=-1)

55
minamo/model/vision.py
View File

@ -1,57 +1,14 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from shared.attention import CBAM
from torchvision.models import resnet18
class MinamoVisionModel(nn.Module):
def __init__(self, tile_types=32, conv_ch=64, out_dim=512):
def __init__(self, tile_types=32, out_dim=512):
super().__init__()
# 输入 softmax 概率值
self.input_conv = nn.Conv2d(tile_types, conv_ch, 3, padding=1)
# 卷积部分
self.vision_conv = nn.Sequential(
nn.Conv2d(conv_ch, conv_ch*2, 3, padding=1),
nn.BatchNorm2d(conv_ch*2),
CBAM(conv_ch*2),
nn.GELU(),
nn.MaxPool2d(2),
nn.Dropout2d(0.4),
nn.Conv2d(conv_ch*2, conv_ch*4, 3, padding=1),
nn.BatchNorm2d(conv_ch*4),
CBAM(conv_ch*4),
nn.GELU(),
nn.MaxPool2d(2),
nn.Dropout2d(0.4),
nn.Conv2d(conv_ch*4, conv_ch*8, 3, padding=1),
nn.BatchNorm2d(conv_ch*8),
CBAM(conv_ch*8),
nn.GELU(),
# nn.MaxPool2d(2),
# nn.Dropout2d(0.4),
nn.Conv2d(conv_ch*8, conv_ch*8, 3, padding=1),
nn.BatchNorm2d(conv_ch*8),
CBAM(conv_ch*8),
nn.GELU(),
nn.AdaptiveMaxPool2d(1)
)
# 输出为向量
self.vision_head = nn.Sequential(
nn.Dropout(0.4),
nn.Linear(conv_ch*8, out_dim)
)
def forward(self, map):
x = self.input_conv(map)
x = self.vision_conv(x)
x = x.view(x.size(0), -1) # 展平
vision_vec = self.vision_head(x)
self.resnet = resnet18(num_classes=out_dim)
self.resnet.conv1 = nn.Conv2d(tile_types, 64, kernel_size=7, stride=2, padding=3, bias=False)
def forward(self, x):
vision_vec = self.resnet(x)
return F.normalize(vision_vec, p=2, dim=-1) # 归一化

3
requirements.txt
View File

@ -3,4 +3,5 @@ torchvision
torchaudio
tqdm
torch-geometric
transformers
transformers
torch-scatter

17
shared/attention.py
View File

@ -54,3 +54,20 @@ class CBAM(nn.Module):
# 空间注意力
s_att = self.spatial_att(x)
return x * s_att
class SEBlock(nn.Module):
def __init__(self, channel, reduction=4):
super().__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(channel, channel // reduction),
nn.GELU(),
nn.Linear(channel // reduction, channel),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.avg_pool(x).view(b, c)
y = self.fc(y).view(b, c, 1, 1)
return x * y

40
shared/graph.py
View File

@ -1,24 +1,23 @@
import torch
from torch_geometric.data import Data, Batch
def differentiable_convert_to_data(map_probs: torch.Tensor) -> Data:
"""
可导的图结构转换（返回PyG Data对象）
可导的图结构转换（返回 PyG Data 对象）
map_probs: [C, H, W]
返回：
Data(x=[N,C], edge_index=[2,E], edge_attr=[E,C])
Data(x=[N, C], edge_index=[2, E], edge_attr=[E, C])
"""
C, H, W = map_probs.shape
device = map_probs.device
N = H * W
# 1. 节点特征（保留所有节点）
# 1. 节点特征
node_features = map_probs.view(C, -1).T # [N, C]
# 2. 构建所有可能的边连接（预计算）
# 生成坐标网格
rows, cols = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
node_indices = rows * W + cols
# 2. 构建所有可能的边连接
node_indices = torch.arange(N, device=device).view(H, W)
# 水平连接（右邻居）
right_src = node_indices[:, :-1].flatten()
@ -28,20 +27,23 @@ def differentiable_convert_to_data(map_probs: torch.Tensor) -> Data:
down_src = node_indices[:-1, :].flatten()
down_dst = node_indices[1:, :].flatten()
# 合并边列表
edge_src = torch.cat([right_src, down_src]).to(device)
edge_dst = torch.cat([right_dst, down_dst]).to(device)
edge_index = torch.stack([edge_src, edge_dst]) # [2, E]
# 合并边列表（双向）
edge_src = torch.cat([right_src, down_src])
edge_dst = torch.cat([right_dst, down_dst])
edge_index = torch.cat([
torch.stack([edge_src, edge_dst], dim=0),
torch.stack([edge_dst, edge_src], dim=0) # 反向连接
], dim=1).to(device, dtype=torch.long)
# 3. 计算可导的边权重（排除墙类型）
wall_class_idx = 1 # 假设类型1是墙
src_probs = 1.0 - map_probs[wall_class_idx].flatten()[edge_src] # [E]
dst_probs = 1.0 - map_probs[wall_class_idx].flatten()[edge_dst] # [E]
edge_mask = (src_probs * dst_probs).unsqueeze(1) # [E, 1]
# 3. 计算可导的边权重
wall_class_idx = 1 # 假设类别 1 是墙
src_probs = torch.sigmoid(-map_probs[wall_class_idx].flatten()[edge_src])
dst_probs = torch.sigmoid(-map_probs[wall_class_idx].flatten()[edge_dst])
edge_mask = torch.nn.functional.softplus(src_probs * dst_probs).unsqueeze(1) # [E, 1]
# 4. 边特征计算（保持可导）
src_feat = map_probs[:, edge_src//W, edge_src%W].T # [E, C]
dst_feat = map_probs[:, edge_dst//W, edge_dst%W].T # [E, C]
# 4. 计算边特征
src_feat = map_probs[:, edge_src // W, edge_src % W].T # [E, C]
dst_feat = map_probs[:, edge_dst // W, edge_dst % W].T # [E, C]
edge_attr = (src_feat + dst_feat) / 2 * edge_mask # [E, C]
return Data(