File size: 3,114 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 | import torch
import torch.nn as nn
import torch_geometric.nn as nng
from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
from onescience.modules.mlp.MLP import StandardMLP
class Model(nn.Module):
"""
PointNet 模型。
用于处理点云数据，通过 MLP 提取局部特征，并使用全局最大池化提取全局特征。
"""
def __init__(self, args, device):
super(Model, self).__init__()
self.__name__ = "PointNet"
# 1. Input Block
self.in_block = StandardMLP(
input_dim=args.n_hidden,
output_dim=args.n_hidden * 2,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
# 2. Max Pooling Block
self.max_block = StandardMLP(
input_dim=args.n_hidden * 2,
output_dim=args.n_hidden * 32,
hidden_dims=[args.n_hidden * 8],
activation=args.act,
use_bias=True
)
# 3. Output Block
self.out_block = StandardMLP(
input_dim=args.n_hidden * (2 + 32), # 34 * hidden
output_dim=args.n_hidden * 4,
hidden_dims=[args.n_hidden * 16],
activation=args.act,
use_bias=True
)
# 4. Encoder
self.encoder = StandardMLP(
input_dim=args.fun_dim + args.space_dim,
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
# 5. Decoder
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
)
self.fcfinal = nn.Linear(args.n_hidden * 4, args.n_hidden)
def forward(self, x, fx, T=None, geo=None):
if geo is None:
raise ValueError("Please provide edge index for Graph Neural Networks")
# 兼容 batch_size = 1 输入
if x.dim() == 3:
x = x.squeeze(0) # [1, N, C] → [N, C]
if fx is not None and fx.dim() == 3:
fx = fx.squeeze(0)
assert (
x.size(0) > 0 # Simple check
), "Input cannot be empty"
# 构造 batch 索引
batch = torch.zeros(x.shape[0], dtype=torch.long, device=x.device)
# 编码 + 局部特征提取
z = torch.cat((x, fx), dim=-1).float()
z = self.encoder(z)
z = self.in_block(z)
# 全局特征（max pooling）
global_coef = self.max_block(z)
global_coef = nng.global_max_pool(global_coef, batch=batch)
# 重复 global coef 到每个点
nb_points = torch.tensor([batch.shape[0]], device=z.device)
global_coef = global_coef.repeat_interleave(nb_points, dim=0)
# 拼接全局 + 局部特征
z = torch.cat([z, global_coef], dim=1)
z = self.out_block(z)
z = self.fcfinal(z)
z = self.decoder(z)
return z.unsqueeze(0)
|