import torch
import torch.nn as nn
from torch.autograd import grad
import numpy as np
import os
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# ================= 工具函数 =================
class AverageMeter(object):
def __init__(self):
self.reset()
def reset(self):
self.val = 0;
self.avg = 0;
self.sum = 0;
self.count = 0
def update(self, val, n=1):
self.val = val;
self.sum += val * n;
self.count += n;
self.avg = self.sum / self.count
def calculate_metrics(y_true, y_pred):
"""计算 MAE, MSE, RMSE, R2"""
y_true = np.array(y_true).flatten()
y_pred = np.array(y_pred).flatten()
mae = mean_absolute_error(y_true, y_pred)
mse = mean_squared_error(y_true, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_true, y_pred)
return mae, mse, rmse, r2
# ================= 基础层 =================
class Sin(nn.Module):
def forward(self, x):
return torch.sin(x)
class MLP(nn.Module):
def __init__(self, input_dim, output_dim, layers_num=4, hidden_dim=50, dropout=0.0):
super(MLP, self).__init__()
self.layers = []
for i in range(layers_num):
in_d = input_dim if i == 0 else hidden_dim
out_d = output_dim if i == layers_num - 1 else hidden_dim
self.layers.append(nn.Linear(in_d, out_d))
if i < layers_num - 1:
self.layers.append(Sin())
if dropout > 0:
self.layers.append(nn.Dropout(p=dropout))
self.net = nn.Sequential(*self.layers)
self._init_weights()
def _init_weights(self):
for layer in self.net:
if isinstance(layer, nn.Linear):
nn.init.xavier_normal_(layer.weight)
def forward(self, x):
return self.net(x)
# ================= 注意力融合模块 =================
class CrossAttentionFusion(nn.Module):
def __init__(self, feature_dim=32, num_heads=2, dropout=0.1):
super().__init__()
self.attn = nn.MultiheadAttention(embed_dim=feature_dim, num_heads=num_heads, batch_first=True, dropout=dropout)
self.norm = nn.LayerNorm(feature_dim)
self.dropout = nn.Dropout(dropout)
# 可学习门控
self.gate_net = nn.Sequential(
nn.Linear(feature_dim * 2, feature_dim),
nn.Tanh(),
nn.Linear(feature_dim, feature_dim),
nn.Sigmoid()
)
def forward(self, curr_feature, hist_feature):
# Query: [B, 1, D], Key/Val: [B, Seq, D]
query = curr_feature.unsqueeze(1)
attn_out, attn_weights = self.attn(query, hist_feature, hist_feature)
attn_out = self.dropout(attn_out.squeeze(1))
# 门控融合
combined = torch.cat([curr_feature, attn_out], dim=1)
gate = self.gate_net(combined)
out = self.norm(curr_feature + gate * attn_out)
return out, attn_weights
# ================= 预测网络 U (增强版) =================
class Solution_u_Enhanced(nn.Module):
def __init__(self, input_dim=17, hidden_dim=60, feature_dim=32):
super(Solution_u_Enhanced, self).__init__()
# 假设 input_dim 是 17 (xt特征)
self.curr_encoder = MLP(input_dim, feature_dim, layers_num=3, hidden_dim=hidden_dim)
self.hist_encoder = MLP(input_dim, feature_dim, layers_num=3, hidden_dim=hidden_dim)
self.fusion = CrossAttentionFusion(feature_dim=feature_dim)
self.predictor = nn.Sequential(
nn.Linear(feature_dim, hidden_dim),
Sin(),
nn.Linear(hidden_dim, 1)
)
def forward(self, xt, history_data=None):
curr_feat = self.curr_encoder(xt)
attn_weights = None
if history_data is not None:
# history: [B, Seq, 17]
b, s, d = history_data.shape
hist_feat = self.hist_encoder(history_data.reshape(-1, d)).reshape(b, s, -1)
fused_feat, attn_weights = self.fusion(curr_feat, hist_feat)
out = self.predictor(fused_feat)
else:
out = self.predictor(curr_feat)
return out, attn_weights
# ================= 物理网络 F (纯净版) =================
class Dynamical_F_Pure(nn.Module):
def __init__(self, input_dim_x=17, hidden_dim=60):
super(Dynamical_F_Pure, self).__init__()
# 输入: x(17) + t(1) + u(1) + u_x(17) + u_t(1) = 37 (假设)
# 这里动态计算维度：input_dim_x + u(1) + u_x(input_dim_x) + u_t(1)
# 注意：通常 PINN 输入是 (x, t)，如果 xt 混合了，维度需根据实际情况调整
# 这里假设传入 [xt, u, u_x, u_t] -> 17 + 1 + 17 + 1 = 36
total_dim = input_dim_x * 2 + 2
self.net = MLP(input_dim=total_dim, output_dim=1, layers_num=3, hidden_dim=hidden_dim)
def forward(self, inputs):
return self.net(inputs)
# ================= 主模型 =================
class AttentionEnhancedPINN(nn.Module):
def __init__(self, args):
super().__init__()
self.args = args
self.input_dim = 17 # 假设输入维度
self.solution_u = Solution_u_Enhanced(input_dim=self.input_dim, hidden_dim=60).to(device)
self.dynamical_F = Dynamical_F_Pure(input_dim_x=self.input_dim, hidden_dim=args.F_hidden_dim).to(device)
# 优化器
self.optimizer = torch.optim.Adam(self.parameters(), lr=args.lr)
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=args.epochs,
eta_min=args.final_lr)
def forward(self, xt, history_data=None):
xt.requires_grad = True
# 1. 预测 u
u, attn_weights = self.solution_u(xt, history_data)
# 2. 自动微分
# 假设 xt 的最后一位是 t，或者 t 隐含在 xt 中
# 这里假设全量求导
grads = grad(u.sum(), xt, create_graph=True, retain_graph=True)[0]
u_t = grads[:, -1:] # 假设最后一列是时间导数
u_x = grads[:, :-1] # 假设前面是空间导数
# 3. 物理残差
f_input = torch.cat([xt, u, u_x, u_t], dim=1)
F_val = self.dynamical_F(f_input)
f = u_t - F_val
return u, f, attn_weights
def train_one_epoch(self, epoch, dataloader, history_loader):
self.train()
loss_meter = AverageMeter()
if self.scheduler: self.scheduler.step()
for batch in dataloader:
# 解析数据：IndexedDataset 返回 (x1, x2, y1, y2, idx1, idx2)
if len(batch) == 6:
x1, x2, y1, y2, idx1, idx2 = batch
hist1 = history_loader.get_batch(idx1)
hist2 = history_loader.get_batch(idx2)
else:
continue
x1, x2, y1, y2 = x1.to(device), x2.to(device), y1.to(device), y2.to(device)
# Forward
u1, f1, attn1 = self.forward(x1, hist1)
u2, f2, _ = self.forward(x2, hist2)
# Losses
l_data = nn.L1Loss()(u1, y1) + nn.L1Loss()(u2, y2) # MAE Loss
l_pde = torch.mean(f1 ** 2) + torch.mean(f2 ** 2) # MSE Loss
l_phy = torch.mean(((u2 - u1) - (y2 - y1)) ** 2) # Consistency
# Regularization
l_reg = 0
if attn1 is not None: l_reg = torch.mean(attn1 ** 2)
# Dynamic Weights
alpha = self.args.alpha if epoch > 20 else 0.0
beta = self.args.beta
loss = l_data + alpha * l_pde + beta * l_phy + self.args.gamma * l_reg
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0)
self.optimizer.step()
loss_meter.update(loss.item())
return loss_meter.avg
def predict(self, loader, history_loader):
"""通用预测函数"""
self.eval()
true_list, pred_list = [], []
with torch.no_grad():
for batch in loader:
if len(batch) >= 5: # 有索引
x = batch[0].to(device)
y = batch[2].to(device)
idx = batch[-1] # 假设最后一个是索引
hist = history_loader.get_batch(idx)
else:
x = batch[0].to(device)
y = batch[2].to(device)
hist = None
u, _ = self.solution_u(x, hist)
true_list.append(y.cpu().numpy())
pred_list.append(u.cpu().numpy())
return np.concatenate(true_list).flatten(), np.concatenate(pred_list).flatten()
def Train(self, trainloader, validloader, testloader, hist_loaders):
train_h, valid_h, test_h = hist_loaders
best_mae = 10.0
for e in range(1, self.args.epochs + 1):
train_loss = self.train_one_epoch(e, trainloader, train_h)
if e % 10 == 0:
# 验证
y_true_v, y_pred_v = self.predict(validloader, valid_h)
mae_v, mse_v, rmse_v, r2_v = calculate_metrics(y_true_v, y_pred_v)
print(f"[Epoch {e}] Loss: {train_loss:.5f} | Valid MAE: {mae_v:.5f} RMSE: {rmse_v:.5f} R2: {r2_v:.4f}")
# 保存最佳并测试
if mae_v < best_mae:
best_mae = mae_v
if testloader:
y_true_t, y_pred_t = self.predict(testloader, test_h)
t_mae, t_mse, t_rmse, t_r2 = calculate_metrics(y_true_t, y_pred_t)
print(
f" >>> Best Model Found! Test Metrics -> MAE: {t_mae:.5f}, MSE: {t_mse:.6f}, RMSE: {t_rmse:.5f}, R2: {t_r2:.4f}")
if self.args.save_folder:
if not os.path.exists(self.args.save_folder): os.makedirs(self.args.save_folder)
np.save(os.path.join(self.args.save_folder, 'true_label.npy'), y_true_t)
np.save(os.path.join(self.args.save_folder, 'pred_label.npy'), y_pred_t)
torch.save(self.state_dict(), os.path.join(self.args.save_folder, 'model.pth'))
我在pinn的基础上加入了注意机制
最新发布
本文介绍如何解析Caffe2的二进制模型，包括init_net和predict_net，通过自定义Python脚本实现对特定层的读取与保存。文章详细解释了proto消息的使用，以及如何通过NetDef结构读取并处理模型数据。
扫一扫
3828
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
