File size: 6,790 Bytes
ff0fadf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 |
import torch
import torch.nn as nn
import torch_geometric.nn as nng
import random
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.sample.SpatialGraphDownsample import SpatialGraphDownsample
from onescience.modules.sample.SpatialGraphUpsample import SpatialGraphUpsample
class Model(nn.Module):
"""
Graph U-Net 模型。
基于图神经网络的 U-Net 结构，包含 Encoder-Decoder 和 Skip Connections。
使用 SpatialGraphDownsample 进行图池化，使用 SpatialGraphUpsample 进行反池化。
"""
def __init__(
self,
args,
device,
pool="random",
scale=5,
list_r=[0.05, 0.2, 0.5, 1, 10],
pool_ratio=[0.5, 0.5, 0.5, 0.5, 0.5],
max_neighbors=64,
layer="SAGE",
head=2,
):
super(Model, self).__init__()
self.__name__ = "Graph_UNet"
# 参数绑定
self.L = scale
self.layer = layer
self.pool_type = pool
self.pool_ratio = pool_ratio
self.list_r = list_r
self.size_hidden = args.n_hidden
self.dim_enc = args.n_hidden
self.bn_bool = True
self.res = False
self.head = head
self.activation = nn.ReLU()
self.encoder = StandardMLP(
input_dim=args.fun_dim,
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
self.decoder = StandardMLP(
input_dim=args.n_hidden,
output_dim=args.out_dim,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
# Down Path Layers
self.down_convs = nn.ModuleList()
self.down_samples = nn.ModuleList()
self.down_bns = nn.ModuleList()
# Level 0 (Initial)
self._add_conv_layer(self.down_convs, self.dim_enc, self.size_hidden)
if self.bn_bool:
self._add_bn_layer(self.down_bns, self.size_hidden)
# Level 1 to L-1
current_dim = self.size_hidden
for n in range(self.L - 1):
self.down_samples.append(
SpatialGraphDownsample(
in_channels=current_dim,
ratio=self.pool_ratio[n],
r=self.list_r[n],
max_num_neighbors=max_neighbors,
pool_method=self.pool_type
)
)
# Conv Layer
in_c = current_dim
out_c = 2 * current_dim if layer == "SAGE" else current_dim
self._add_conv_layer(self.down_convs, in_c, out_c)
current_dim = out_c
if self.bn_bool:
self._add_bn_layer(self.down_bns, current_dim)
# Up Path Layers
self.up_convs = nn.ModuleList()
# --- 3. Upsample Module ---
self.up_sampler = SpatialGraphUpsample()
self.up_bns = nn.ModuleList()
curr_h_init = args.n_hidden
# Up Layer 0 (Top Layer)
if self.layer == "SAGE":
self.up_convs.append(nng.SAGEConv(3 * curr_h_init, self.dim_enc))
curr_h_init = 2 * curr_h_init
elif self.layer == "GAT":
self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, self.dim_enc, heads=2, concat=False))
if self.bn_bool:
self.up_bns.append(nng.BatchNorm(self.dim_enc, track_running_stats=False))
# Up Layer 1 to L-1 (Middle Layers)
for n in range(1, self.L - 1):
if self.layer == "SAGE":
self.up_convs.append(nng.SAGEConv(3 * curr_h_init, curr_h_init))
bn_dim = curr_h_init
curr_h_init = 2 * curr_h_init
elif self.layer == "GAT":
self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, curr_h_init, heads=2, concat=True))
bn_dim = curr_h_init * 2 # GAT concat=True
if self.bn_bool:
self.up_bns.append(nng.BatchNorm(bn_dim, track_running_stats=False))
def _add_conv_layer(self, module_list, in_c, out_c):
if self.layer == "SAGE":
module_list.append(nng.SAGEConv(in_c, out_c))
elif self.layer == "GAT":
module_list.append(nng.GATConv(in_c, out_c, heads=self.head, concat=True, add_self_loops=False))
def _add_bn_layer(self, module_list, in_c):
dim = in_c * self.head if self.layer == "GAT" else in_c
module_list.append(nng.BatchNorm(dim, track_running_stats=False))
def forward(self, x, fx, T=None, geo=None):
if geo is None: raise ValueError("Edge index required")
if fx.dim() == 3: fx = fx.squeeze(0)
if geo.dim() == 3: edge_index = geo.squeeze(0)
else: edge_index = geo
# Encoder
z = self.encoder(fx)
if self.res: z_res = z.clone()
# Downsampling Path
skip_connections = []
pos_history = []
edge_index_history = [edge_index.clone()]
# Level 0 Conv
z = self.down_convs[0](z, edge_index)
if self.bn_bool: z = self.down_bns[0](z)
z = self.activation(z)
skip_connections.append(z.clone())
# Assuming x contains coords in first 2 columns as per original code logic
current_pos = x[:, :2]
pos_history.append(current_pos.clone())
# Levels 1 to L-1
for n in range(self.L - 1):
z, current_pos, edge_index, _ = self.down_samples[n](z, current_pos, edge_index)
pos_history.append(current_pos.clone())
edge_index_history.append(edge_index.clone())
z = self.down_convs[n+1](z, edge_index)
if self.bn_bool: z = self.down_bns[n+1](z)
z = self.activation(z)
skip_connections.append(z.clone())
# Up Path
for n in range(self.L - 1, 0, -1):
layer_idx = n - 1
pos_low = pos_history[n]
pos_high = pos_history[n-1]
z_skip = skip_connections[n-1]
target_edge_index = edge_index_history[n-1]
z = self.up_sampler(z, pos_low, pos_high)
z = torch.cat([z, z_skip], dim=1)
z = self.up_convs[layer_idx](z, target_edge_index)
if self.bn_bool:
z = self.up_bns[layer_idx](z)
if n != 1:
z = self.activation(z)
# Decoder
if self.res: z = z + z_res
z = self.decoder(z)
return z.unsqueeze(0)
|