技巧:
比如下面一个新的模型deeponet,我自己的数据很复杂,这里在代码最后用用随机生成的数据,两分钟就完成了代码的测试成功。
import torch
import torch.nn as nn
import torch.optim as optim# 带偏置项的 DeepONet 结构,包括 Branch 和 Trunk 网络
class DeepONet(nn.Module):def __init__(self, branch_input_dim, trunk_input_dim, hidden_dim):super(DeepONet, self).__init__()# Branch 网络,用于处理输入点云的特征(例如位移量、压强)self.branch_net = nn.Sequential(nn.Linear(branch_input_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim))# Trunk 网络,用于处理时间和空间坐标 [x, y, z, t]self.trunk_net = nn.Sequential(nn.Linear(trunk_input_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim))# 偏置项 biasself.bias = nn.Parameter(torch.zeros(1)) # 可训练的偏置项# 最终的输出层,预测位移或压强等物理状态self.fc_output = nn.Linear(hidden_dim, 3)def forward(self, point_features, coord_time):# Branch网络的输出branch_output = self.branch_net(point_features)# Trunk网络的输出trunk_output = self.trunk_net(coord_time)# 将 Branch 和 Trunk 的输出结合,计算最终的输出combined = branch_output * trunk_outputoutput = self.fc_output(combined) + self.bias # 加上偏置项return output# 数据准备
# 输入的数据格式:
# point_features:3D点云的物理特征(例如位移量 pointDisplacement、压强 p)
# coord_time:空间位置和时间 [x, y, z, t]# 示例数据的维度设置
branch_input_dim = 3 # 例如 [pointDisplacement, p, ...]
trunk_input_dim = 4 # [x, y, z, t]
hidden_dim = 64 # 隐藏层维度,可根据需求调整# 模型初始化
model = DeepONet(branch_input_dim, trunk_input_dim, hidden_dim)# 损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练流程
def train(model, point_features, coord_time, target, epochs=1000):for epoch in range(epochs):optimizer.zero_grad()# 前向传播output = model(point_features, coord_time)# 计算损失loss = criterion(output, target)# 反向传播和优化loss.backward()optimizer.step()if epoch % 100 == 0:print(f"Epoch {epoch}, Loss: {loss.item()}")# 示例数据,实际应用时需要替换为真实数据
N = 1000 # 样本数量
point_features = torch.randn(N, branch_input_dim) # 3D点云的物理特征
coord_time = torch.randn(N, trunk_input_dim) # [x, y, z, t]
target = torch.randn(N, 3) # 目标物理状态# 训练模型
train(model, point_features, coord_time, target, epochs=1000)# 推理:给定新的时空点,预测物理状态
def predict(model, point_features, coord_time):model.eval()with torch.no_grad():prediction = model(point_features, coord_time)return prediction# 示例推理
new_point_features = torch.randn(1, branch_input_dim)
new_coord_time = torch.tensor([[0.5, 0.5, 0.5, 0.1]]) # 在 t=0.1 的 (0.5, 0.5, 0.5) 空间点
prediction = predict(model, new_point_features, new_coord_time)
print("Predicted state:", prediction)
输出如下:
Epoch 0, Loss: 1.0260347127914429
Epoch 100, Loss: 0.7669863104820251
Epoch 200, Loss: 0.5786211490631104
Epoch 300, Loss: 0.4749055504798889
Epoch 400, Loss: 0.41076529026031494
Epoch 500, Loss: 0.36538082361221313
Epoch 600, Loss: 0.39494913816452026
Epoch 700, Loss: 0.30206459760665894
Epoch 800, Loss: 0.2839098572731018
Epoch 900, Loss: 0.2648167908191681
Predicted state: tensor([[-0.2604, 0.2214, 0.5066]])Process finished with exit code 0