feat: maskGIT 加入热力图条件限定

ginka/dataset.py

@@ -237,4 +237,26 @@ class GinkaRNNDataset(Dataset):
             "tag_cond": tag_cond,
             "val_cond": val_cond,
             "target_map": target
-        }
+        }
+        
+class GinkaMaskGITDataset(Dataset):
+    def __init__(self, data_path: str, device):
+        self.data = load_data(data_path)  # 自定义数据加载函数
+        self.device = device
+        
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        
+        target = torch.LongTensor(item['map']) # [H, W]
+        cond = torch.FloatTensor(item['val']) # [cond_dim]
+        heatmap = torch.FloatTensor(item['heatmap']) # [heatmap_channel, H, W]
+
+        return {
+            "cond": cond,
+            "target_map": target,
+            "heatmap": heatmap
+        }
+        

ginka/maskGIT/cond.py

@@ -0,0 +1,58 @@
+import time
+import torch
+import torch.nn as nn
+from ..utils import print_memory
+
+class GinkaMaskGITCond(nn.Module):
+    def __init__(self, cond_dim=16, heatmap_channel=4, output_dim=256):
+        super().__init__()
+        self.cond_fc = nn.Sequential(
+            nn.Linear(cond_dim, output_dim // 2),
+            nn.LayerNorm(output_dim // 2),
+            nn.ReLU(),
+            
+            nn.Linear(output_dim // 2, output_dim)
+        )
+        
+        self.heatmap_conv = nn.Sequential(
+            nn.Conv2d(heatmap_channel, output_dim // 4, kernel_size=3, padding=1, padding_mode='replicate'),
+            nn.BatchNorm2d(output_dim // 4),
+            nn.ReLU(),
+            
+            nn.Conv2d(output_dim // 4, output_dim // 2, kernel_size=3, padding=1, padding_mode='replicate'),
+            nn.BatchNorm2d(output_dim // 2),
+            nn.ReLU(),
+            
+            nn.Conv2d(output_dim // 2, output_dim, kernel_size=3, padding=1, padding_mode='replicate')
+        )
+    
+    def forward(self, cond, heatmap):
+        # cond: [B, cond_dim]
+        # heatmap: [B, C, H, W]
+        cond = self.cond_fc(cond)
+        heatmap = self.heatmap_conv(heatmap)
+        return cond, heatmap
+    
+if __name__ == "__main__":
+    device = torch.device("cpu")
+    
+    cond = torch.rand(1, 16).to(device)
+    heatmap = torch.rand(1, 4, 13, 13).to(device)
+    
+    # 初始化模型
+    model = GinkaMaskGITCond().to(device)
+    
+    print_memory("初始化后")
+    
+    # 前向传播
+    start = time.perf_counter()
+    cond, heatmap = model(cond, heatmap)
+    end = time.perf_counter()
+    
+    print_memory("前向传播后")
+    
+    print(f"推理耗时: {end - start}")
+    print(f"输出形状: cond={cond.shape}, heatmap={heatmap.shape}")
+    print(f"Cond FC parameters: {sum(p.numel() for p in model.cond_fc.parameters())}")
+    print(f"Heatmap Conv parameters: {sum(p.numel() for p in model.heatmap_conv.parameters())}")
+    print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")

ginka/maskGIT/model.py

@@ -2,19 +2,19 @@ import time
 import torch
 import torch.nn as nn
 from ..utils import print_memory
+from .cond import GinkaMaskGITCond
 
 class GinkaMaskGIT(nn.Module):
     def __init__(
-        self, num_classes=16, cond_dim=16, d_model=256, dim_ff=512, nhead=8, num_layers=4, map_size=13*13
+        self, num_classes=16, cond_dim=16, heatmap_channel=4, d_model=256, 
+        dim_ff=512, nhead=8, num_layers=4, map_size=13*13
     ):
         super().__init__()
         
         self.tile_embedding = nn.Embedding(num_classes, d_model)
-        self.pos_embedding = nn.Parameter(torch.randn(1, map_size, d_model))
+        self.pos_embedding = nn.Parameter(torch.randn(1, map_size + 1, d_model))
         
-        self.cond_projection = nn.Sequential(
-            nn.Linear(cond_dim, d_model)
-        )
+        self.cond_encoder = GinkaMaskGITCond(cond_dim=cond_dim, heatmap_channel=heatmap_channel, output_dim=d_model)
         
         self.encoder = nn.TransformerEncoder(
             nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, batch_first=True),
@@ -29,14 +29,20 @@ class GinkaMaskGIT(nn.Module):
             nn.Linear(d_model, num_classes)
         )
         
-    def forward(self, map: torch.Tensor, cond: torch.Tensor):
+    def forward(self, map: torch.Tensor, cond: torch.Tensor, heatmap: torch.Tensor):
         # map: [B, H * W]
         # cond: [B, cond_dim]
+        # heatmap: [B, C, H, W]
         # output: [B, H * W, num_classes]
+        cond, heatmap = self.cond_encoder(cond, heatmap)
+        # cond: [B, d_model]
+        # heatmap: [B, d_model, H, W]
         
-        x = self.tile_embedding(map) + self.pos_embedding
-        c = self.cond_projection(cond).unsqueeze(1)
-        x = torch.cat([c, x], dim=1)
+        B, C, H, W = heatmap.shape
+        
+        heatmap = heatmap.view(B, C, H * W).permute(0, 2, 1)
+        x = self.tile_embedding(map) + heatmap
+        x = torch.cat([cond.unsqueeze(1), x], dim=1) + self.pos_embedding
         
         m = self.encoder(x)
         out = self.decoder(x, m)
@@ -50,6 +56,7 @@ if __name__ == "__main__":
     
     map = torch.randint(0, 16, [1, 169]).to(device)
     cond = torch.rand(1, 16).to(device)
+    heatmap = torch.rand(1, 4, 13, 13).to(device)
     
     # 初始化模型
     model = GinkaMaskGIT().to(device)
@@ -58,7 +65,7 @@ if __name__ == "__main__":
     
     # 前向传播
     start = time.perf_counter()
-    output = model(map, cond)
+    output = model(map, cond, heatmap)
     end = time.perf_counter()
     
     print_memory("前向传播后")
@@ -66,7 +73,7 @@ if __name__ == "__main__":
     print(f"推理耗时: {end - start}")
     print(f"输出形状: output={output.shape}")
     print(f"Tile Embedding parameters: {sum(p.numel() for p in model.tile_embedding.parameters())}")
-    print(f"Projection parameters: {sum(p.numel() for p in model.cond_projection.parameters())}")
+    print(f"Condition Encoder parameters: {sum(p.numel() for p in model.cond_encoder.parameters())}")
     print(f"Encoder parameters: {sum(p.numel() for p in model.encoder.parameters())}")
     print(f"Decoder parameters: {sum(p.numel() for p in model.decoder.parameters())}")
     print(f"Output parameters: {sum(p.numel() for p in model.output_fc.parameters())}")

ginka/train_maskGIT.py

@@ -11,12 +11,10 @@ import cv2
 import numpy as np
 from torch_geometric.loader import DataLoader
 from tqdm import tqdm
-from .transformer.maskGIT import GinkaMaskGIT
-from .vae_rnn.loss import VAELoss
-from .vae_rnn.scheduler import VAEScheduler
-from .dataset import GinkaRNNDataset
+from .maskGIT.model import GinkaMaskGIT
+from .dataset import GinkaMaskGITDataset
 from shared.image import matrix_to_image_cv
-from .transformer.mask import MapMask
+from .maskGIT.mask import MapMask
 
 # 手工标注标签定义（暂时不用）：
 # 0. 蓝海, 1. 红海, 2: 室内, 3. 野外, 4. 左右对称, 5. 上下对称, 6. 伪对称, 7. 咸鱼层,
@@ -83,8 +81,8 @@ def train():
     model = GinkaMaskGIT(num_classes=NUM_CLASSES).to(device)
     masker = MapMask([0.5, 0.5])
     
-    dataset = GinkaRNNDataset(args.train, device)
-    dataset_val = GinkaRNNDataset(args.validate, device)
+    dataset = GinkaMaskGITDataset(args.train, device)
+    dataset_val = GinkaMaskGITDataset(args.validate, device)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
     dataloader_val = DataLoader(dataset_val, batch_size=BATCH_SIZE // VAL_BATCH_DIVIDER, shuffle=True)