import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import time
import gc
# 1. 定义物理信息神经网络
class SenseNet(nn.Module):
def __init__(self, input_dim, output_dim, hidden_layers=[20, 30, 30, 20]):
super().__init__()
layers = []
layers.append(nn.Linear(input_dim, hidden_layers[0]))
layers.append(nn.Tanh())
for i in range(len(hidden_layers)-1):
layers.append(nn.Linear(hidden_layers[i], hidden_layers[i+1]))
layers.append(nn.Tanh())
layers.append(nn.Linear(hidden_layers[-1], output_dim))
self.net = nn.Sequential(*layers)
# 材料参数 (线性弹性)
self.E = 100.0 # 弹性模量
self.nu = 0.4 # 泊松比
self.lam = self.E * self.nu / ((1 + self.nu) * (1 - 2 * self.nu))
self.mu = self.E / (2 * (1 + self.nu))
def forward(self, x):
return self.net(x)
def compute_strain(self, x):
"""通过自动微分计算应变张量(优化版本)"""
# 确保输入需要梯度
if not x.requires_grad:
x = x.clone().requires_grad_(True)
# 计算位移
u = self.net(x)
# 初始化雅可比矩阵
jac = torch.zeros(x.size(0), u.size(1), x.size(1), device=x.device)
# 为每个输出分量计算梯度
for i in range(u.size(1)):
# 创建梯度输出
grad_outputs = torch.ones(u.size(0), device=x.device)
# 计算梯度
gradients = torch.autograd.grad(
outputs=u[:, i],
inputs=x,
grad_outputs=grad_outputs,
create_graph=True,
retain_graph=True,
only_inputs=True,
allow_unused=True # 关键修改:允许未使用变量
)[0]
if gradients is not None:
jac[:, i, :] = gradients
# 计算应变张量 (对称部分)
strain = 0.5 * (jac + jac.transpose(1, 2))
return strain
def compute_stress(self, strain):
"""通过本构关系计算应力张量"""
# 计算应变迹
trace_e = torch.diagonal(strain, dim1=1, dim2=2).sum(dim=1)
# 创建单位张量
identity = torch.eye(strain.shape[2], device=strain.device)
identity = identity.unsqueeze(0).repeat(strain.shape[0], 1, 1)
# 计算应力张量
stress = self.lam * trace_e[:, None, None] * identity + 2 * self.mu * strain
return stress
def compute_equilibrium_residual(self, x):
"""计算平衡方程残差 ∇·σ = 0"""
# 使用更高效的方法计算散度
strain = self.compute_strain(x)
stress = self.compute_stress(strain)
# 计算应力张量的散度
div_stress = torch.zeros(x.size(0), x.size(1), device=x.device)
for i in range(x.size(1)):
# 对每个方向计算散度
div = torch.zeros(x.size(0), device=x.device)
for j in range(stress.size(1)):
grad = torch.autograd.grad(
outputs=stress[:, j, i].sum(),
inputs=x,
create_graph=True,
retain_graph=True,
only_inputs=True,
allow_unused=True, # 关键修改:允许未使用变量
grad_outputs=torch.ones_like(stress[:, j, i].sum())
)[0]
if grad is not None:
div += grad[:, j]
div_stress[:, i] = div
return div_stress
def loss_function(self, x, sensor_data):
"""组合损失函数"""
# 边界条件损失 (示例:左侧固定)
left_boundary_mask = x[:, 0] == 0.0
left_boundary = x[left_boundary_mask]
if len(left_boundary) > 0:
u_left = self.net(left_boundary)
mse_u = torch.mean(u_left[:, 0]**2 + u_left[:, 1]**2) # 假设全固定
else:
mse_u = torch.tensor(0.0, device=x.device)
# 应变匹配损失
sensor_points = sensor_data['locations']
sensor_strain = sensor_data['values']
# 确保传感器点需要梯度
if not sensor_points.requires_grad:
sensor_points = sensor_points.clone().requires_grad_(True)
# 仅计算传感器点的应变
strain_pred = self.compute_strain(sensor_points)
# 转换为传感器方向的投影应变 (简化处理)
projected_strain = strain_pred[:, 0, 0] # 假设测量x方向应变
mse_e = torch.mean((projected_strain - sensor_strain)**2)
# 平衡方程损失 (使用随机采样点减少计算量)
sample_indices = torch.randperm(x.size(0))[:500] # 随机采样500点
x_sample = x[sample_indices]
# 关键修改:分离计算图以避免反向传播问题
with torch.no_grad():
residual = self.compute_equilibrium_residual(x_sample)
mse_f = torch.mean(torch.sum(residual**2, dim=1))
# 加权组合损失 (根据论文设置)
total_loss = 10000 * mse_u + 10000 * mse_e + 1 * mse_f
return total_loss
# 2. 生成模拟数据 (替代有限元仿真)
def generate_synthetic_data(num_points=1000):
"""生成弯曲变形的合成数据"""
print("Generating synthetic data...")
# 创建坐标网格
grid_size = int(np.sqrt(num_points))
x = torch.linspace(0, 5, grid_size)
y = torch.linspace(0, 1, grid_size)
# 修复meshgrid警告
X, Y = torch.meshgrid(x, y, indexing='ij')
# 创建坐标张量 (保留梯度)
coords = torch.stack([X.ravel(), Y.ravel()], dim=1).float().requires_grad_(True)
# 生成位移场 (弯曲变形)
ux = 0.01 * (X**2) * Y
uy = 0.02 * X * (Y**2)
displacement = torch.stack([ux.ravel(), uy.ravel()], dim=1)
# 生成应变数据 (仅表面)
sensor_indices = torch.randperm(coords.size(0))[:grid_size] # 随机选择传感器点
sensor_locations = coords[sensor_indices].clone() # 保留梯度
strain_xx = 0.02 * X * Y # 简化应变计算
sensor_strain = strain_xx.ravel()[sensor_indices].detach().clone() # 应变值不需要梯度
print(f"Generated {coords.size(0)} points with {sensor_locations.size(0)} sensors")
return {
'coords': coords,
'displacement': displacement.detach(),
'sensor_data': {
'locations': sensor_locations,
'values': sensor_strain
}
}
# 3. 训练过程
def train_sensenet():
# 配置
input_dim = 2
output_dim = 2
epochs = 100 # 减少epoch数以加快训练
lr = 0.001
# 生成数据
print("Generating synthetic data...")
data = generate_synthetic_data(num_points=1000)
coords = data['coords']
sensor_data = data['sensor_data']
# 初始化模型和优化器
print("Initializing model...")
model = SenseNet(input_dim, output_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.995)
# 训练循环
print("Starting training...")
losses = []
start_time = time.time()
# 启用梯度追踪
torch.set_grad_enabled(True)
for epoch in range(epochs):
optimizer.zero_grad()
loss = model.loss_function(coords, sensor_data)
loss.backward()
optimizer.step()
scheduler.step()
losses.append(loss.item())
if epoch % 10 == 0: # 每10个epoch输出一次
elapsed = time.time() - start_time
print(f"Epoch {epoch}/{epochs}, Loss: {loss.item():.4f}, Time: {elapsed:.1f}s")
start_time = time.time()
# 定期清理内存
if epoch % 50 == 0:
gc.collect()
# 可视化结果
print("Visualizing results...")
with torch.no_grad():
# 预测位移
pred = model(coords).detach().numpy()
true = data['displacement'].numpy()
# 计算误差
error = np.sqrt(np.sum((pred - true)**2, axis=1))
max_error = np.max(error)
mean_error = np.mean(error)
# 创建可视化图表
plt.figure(figsize=(18, 12))
# 真实位移
plt.subplot(231)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=true[:, 0], cmap='jet', s=10)
plt.title("True Displacement (ux)")
plt.colorbar()
# 预测位移
plt.subplot(232)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=pred[:, 0], cmap='jet', s=10)
plt.title("Predicted Displacement (ux)")
plt.colorbar()
# 误差
plt.subplot(233)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=pred[:, 0]-true[:, 0], cmap='coolwarm', s=10)
plt.title("Error in ux Prediction")
plt.colorbar()
# 真实位移
plt.subplot(234)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=true[:, 1], cmap='jet', s=10)
plt.title("True Displacement (uy)")
plt.colorbar()
# 预测位移
plt.subplot(235)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=pred[:, 1], cmap='jet', s=10)
plt.title("Predicted Displacement (uy)")
plt.colorbar()
# 误差
plt.subplot(236)
plt.scatter(coords[:, 0].detach().numpy(), coords[:, 1].detach().numpy(),
c=pred[:, 1]-true[:, 1], cmap='coolwarm', s=10)
plt.title("Error in uy Prediction")
plt.colorbar()
plt.suptitle(f"Max Error: {max_error:.4f}, Mean Error: {mean_error:.4f}")
plt.tight_layout()
plt.savefig('displacement_comparison.png')
# 训练损失曲线
plt.figure(figsize=(10, 6))
plt.plot(losses)
plt.title("Training Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.yscale('log')
plt.grid(True)
plt.savefig('training_loss.png')
plt.show()
return model
# 运行训练
if __name__ == "__main__":
print("Starting SenseNet training...")
print(f"PyTorch version: {torch.__version__}")
print(f"Device: {'GPU' if torch.cuda.is_available() else 'CPU'}")
# 设置随机种子以确保可重复性
torch.manual_seed(42)
np.random.seed(42)
trained_model = train_sensenet()
print("Training completed!")这段代码进行应变场重构有什么错误吗
最新发布