技巧:
比如下面一个新的模型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)
)
# 偏置项 bias
self.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_output
output = 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