mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-05-14 04:41:12 +08:00
25 lines
866 B
Python
25 lines
866 B
Python
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from .vision import MinamoVisionModel
|
|
from .topo import MinamoTopoModel
|
|
|
|
class MinamoModel(nn.Module):
|
|
def __init__(self, tile_types=32):
|
|
super().__init__()
|
|
# 视觉相似度部分
|
|
self.vision_model = MinamoVisionModel(tile_types)
|
|
# 拓扑相似度部分
|
|
self.topo_model = MinamoTopoModel(tile_types)
|
|
|
|
def forward(self, map1, map2, graph1, graph2):
|
|
vision_feat1 = self.vision_model(map1)
|
|
vision_feat2 = self.vision_model(map2)
|
|
|
|
topo_feat1 = self.topo_model(graph1)
|
|
topo_feat2 = self.topo_model(graph2)
|
|
|
|
vision_sim = F.cosine_similarity(vision_feat1, vision_feat2, -1).unsqueeze(-1)
|
|
topo_sim = F.cosine_similarity(topo_feat1, topo_feat2, -1).unsqueeze(-1)
|
|
|
|
return vision_sim, topo_sim
|