import os
import csv
import torch
import scipy.io
import numpy as np
import pandas as pd
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as plt
from copy import deepcopy
from itertools import product
from sklearn.metrics import mean_absolute_error, r2_score
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
#####数据加载#####
def load_multi_dun_exp_data(file_path, device, num_samples=None):
"""
从单个.mat文件加载多墩墩数据
"""
print(f"Loading data from: {file_path}")
mat_data = scipy.io.loadmat(file_path)
# 加载力数据 (总合力)
fx = mat_data['fx'].flatten()
fy = mat_data['fy'].flatten()
fz = -mat_data['fz'].flatten() # fz反号
F_total = np.column_stack((fx, fy, fz))
# 加载磁场数据
B_processed = mat_data['B_processed'] # 形状: (N, cols)
# 推断墩墩数量：cols / 18
total_cols = B_processed.shape[1]
num_duns = total_cols // 18
if total_cols % 18 != 0:
raise ValueError(f"B_processed列数({total_cols})不能被18整除")
print(f"Inferred number of duns: {num_duns}")
# 样本选择
total_samples = len(F_total)
if num_samples is None or num_samples > total_samples:
print(f"Using all {total_samples} samples")
indices = np.arange(total_samples)
else:
indices = np.random.choice(total_samples, size=num_samples, replace=False)
indices = np.sort(indices)
print(f"Randomly selected {num_samples} samples from {total_samples}")
# 选择数据
F_total_selected = F_total[indices]
B_processed_selected = B_processed[indices, :]
# 重新组织数据格式: (N, num_duns, 6, 3)
B_array = np.zeros((len(indices), num_duns, 6, 3))
for dun_id in range(num_duns):
col_start = dun_id * 18
col_end = (dun_id + 1) * 18
B_dun = B_processed_selected[:, col_start:col_end] # (N, 18)
B_array[:, dun_id, :, :] = B_dun.reshape(-1, 6, 3)
B_tensor = torch.Tensor(B_array).to(device)
F_tensor = torch.Tensor(F_total_selected).to(device)
print(f"Final data shape - B: {B_tensor.shape}, F: {F_tensor.shape}")
return B_tensor, F_tensor, num_duns
def load_train_data(train_data_path, device):
"""
加载预训练数据（单个墩墩）
"""
train_data = pd.read_csv(train_data_path)
force_cols = ['Fx_total', 'Fy_total', 'Fz_total']
field_cols = [f'B_{k}{j}' for k, j in product(range(1, 7), ['x', 'y', 'z'])]
F_train = torch.Tensor(train_data[force_cols].values).to(device)
B_train = torch.Tensor(train_data[field_cols].values).reshape(-1, 6, 3).to(device)
print(f"Pretrain data - B: {B_train.shape}, F: {F_train.shape}")
return B_train, F_train
#####模型构建#####
class PretrainedEncoder(nn.Module):
"""预训练编码器 - 从单个墩墩模型加载"""
def __init__(self, input_dim=18, hidden_dim=64, output_dim=32, dropout_rate=0.3):
super(PretrainedEncoder, self).__init__()
self.network = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.GELU(),
nn.Dropout(dropout_rate),
nn.Linear(hidden_dim, output_dim),
nn.BatchNorm1d(output_dim),
nn.GELU(),
nn.Dropout(dropout_rate),
)
def forward(self, x):
# x: (batch_size, 6, 3) -> (batch_size, 18)
x = x.reshape(x.size(0), -1)
return self.network(x)
class ArrayForceInversionModel(nn.Module):
"""
阵列力反演模型
架构: 6个预训练编码器 + Transformer交互 + 合力预测器
"""
def __init__(self, num_duns, encoder_dim=32, dropout_rate=0.2):
super(ArrayForceInversionModel, self).__init__()
self.num_duns = num_duns
self.encoder_dim = encoder_dim
# 6个预训练编码器（每个墩墩一个）
self.dun_encoders = nn.ModuleList([
PretrainedEncoder(output_dim=encoder_dim, dropout_rate=dropout_rate)
for _ in range(num_duns)
])
# Transformer交互网络
encoder_layer = nn.TransformerEncoderLayer(
d_model=encoder_dim,
nhead=8,
dim_feedforward=128,
dropout=dropout_rate,
batch_first=True
)
self.interaction_network = nn.TransformerEncoder(encoder_layer, num_layers=3)
# 位置编码
self.position_encoding = nn.Parameter(torch.randn(1, num_duns, encoder_dim))
# 合力预测器
self.force_predictor = nn.Sequential(
nn.Linear(num_duns * encoder_dim, 128),
nn.GELU(),
nn.Dropout(dropout_rate),
nn.Linear(128, 64),
nn.GELU(),
nn.Dropout(dropout_rate),
nn.Linear(64, 3) # 输出三维合力
)
# 归一化参数
self.register_buffer('B_mean', torch.zeros(18))
self.register_buffer('B_std', torch.ones(18))
self.register_buffer('F_mean', torch.zeros(3))
self.register_buffer('F_std', torch.ones(3))
def set_normalization_params(self, B_mean, B_std, F_mean, F_std):
"""设置归一化参数"""
self.B_mean = B_mean.to(device)
self.B_std = B_std.to(device)
self.F_mean = F_mean.to(device)
self.F_std = F_std.to(device)
def normalize_B(self, B):
"""归一化磁场数据"""
batch_size, num_duns, num_magnets, dim = B.shape
B_flat = B.reshape(batch_size, num_duns, -1)
B_norm = (B_flat - self.B_mean) / (self.B_std + 1e-8)
return B_norm.reshape(batch_size, num_duns, num_magnets, dim)
def denormalize_F(self, F):
"""反归一化力数据"""
return F * self.F_std + self.F_mean
def forward(self, B_array, return_features=False):
"""
Args:
B_array: (batch_size, num_duns, 6, 3) - 阵列磁场数据
return_features: 是否返回中间特征
Returns:
F_total: (batch_size, 3) - 预测的总合力
features: (可选) 编码器特征
"""
batch_size = B_array.size(0)
# 归一化输入
B_norm = self.normalize_B(B_array)
# 分别编码每个墩墩
dun_features = []
for dun_id in range(self.num_duns):
B_dun = B_norm[:, dun_id, :, :] # (batch_size, 6, 3)
features = self.dun_encoders[dun_id](B_dun) # (batch_size, encoder_dim)
dun_features.append(features)
# 堆叠特征并添加位置编码
features_tensor = torch.stack(dun_features, dim=1) # (batch_size, num_duns, encoder_dim)
features_tensor = features_tensor + self.position_encoding
# Transformer交互
interacted_features = self.interaction_network(features_tensor) # (batch_size, num_duns, encoder_dim)
# 预测合力
flattened = interacted_features.reshape(batch_size, -1) # (batch_size, num_duns * encoder_dim)
F_pred_norm = self.force_predictor(flattened) # (batch_size, 3)
F_pred = self.denormalize_F(F_pred_norm)
if return_features:
return F_pred, dun_features, interacted_features
return F_pred
#####损失函数#####
class ImprovedMMDLoss(nn.Module):
"""改进的MMD损失，支持多尺度核"""
def __init__(self, kernel_mul=2.0, kernel_num=5):
super(ImprovedMMDLoss, self).__init__()
self.kernel_mul = kernel_mul
self.kernel_num = kernel_num
def guassian_kernel(self, source, target):
n_samples = source.size(0) + target.size(0)
total = torch.cat([source, target], dim=0)
total0 = total.unsqueeze(0).expand(total.size(0), total.size(0), total.size(1))
total1 = total.unsqueeze(1).expand(total.size(0), total.size(0), total.size(1))
L2_distance = ((total0 - total1) ** 2).sum(2)
bandwidth = torch.sum(L2_distance.data) / (n_samples ** 2 - n_samples)
bandwidth /= self.kernel_mul ** (self.kernel_num // 2)
bandwidth_list = [bandwidth * (self.kernel_mul ** i) for i in range(self.kernel_num)]
kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
return sum(kernel_val)
def forward(self, source, target):
batch_size = source.size(0)
kernels = self.guassian_kernel(source, target)
XX = kernels[:batch_size, :batch_size]
YY = kernels[batch_size:, batch_size:]
XY = kernels[:batch_size, batch_size:]
YX = kernels[batch_size:, :batch_size]
loss = torch.mean(XX + YY - XY - YX)
return loss
class PhysicsConsistencyLoss(nn.Module):
"""物理一致性损失"""
def __init__(self, symmetry_pairs=None):
super(PhysicsConsistencyLoss, self).__init__()
self.symmetry_pairs = symmetry_pairs or [(0,5), (1,4), (2,3)] # 默认对称对
def forward(self, dun_features):
"""
Args:
dun_features: list of (batch_size, feature_dim) - 每个墩墩的特征
"""
consistency_loss = 0.0
# 对称性约束
for i, j in self.symmetry_pairs:
if i < len(dun_features) and j < len(dun_features):
# 对称墩墩的特征应该相似
feat_i = dun_features[i]
feat_j = dun_features[j]
consistency_loss += F.mse_loss(feat_i, feat_j)
# 特征平滑性约束
for i in range(len(dun_features) - 1):
consistency_loss += 0.1 * F.mse_loss(dun_features[i], dun_features[i+1])
return consistency_loss
#####评估指标#####
def compute_metrics(F_pred, F_true):
"""计算评估指标"""
F_pred_np = F_pred.cpu().numpy()
F_true_np = F_true.cpu().numpy()
metrics = {}
force_components = ['Fx', 'Fy', 'Fz']
# 各分量指标
for i, comp in enumerate(force_components):
pred_comp = F_pred_np[:, i]
true_comp = F_true_np[:, i]
metrics[comp] = {
'mae': mean_absolute_error(true_comp, pred_comp),
'r2': r2_score(true_comp, pred_comp)
}
# 总体指标
total_mae = mean_absolute_error(F_true_np, F_pred_np)
total_r2 = r2_score(F_true_np, F_pred_np)
metrics['overall'] = {
'mae': total_mae,
'r2': total_r2
}
return metrics
#####微调程序#####
class ArrayForceInversionFineTuner:
def __init__(self, pretrained_path, output_folder, num_duns=6):
self.pretrained_path = pretrained_path
self.output_folder = output_folder
self.num_duns = num_duns
self.model = None
self.history = []
os.makedirs(output_folder, exist_ok=True)
# 损失函数
self.mse_loss = nn.MSELoss()
self.mmd_loss = ImprovedMMDLoss()
self.physics_loss = PhysicsConsistencyLoss()
def load_pretrained_encoders(self, B_train_sample, F_train_sample):
"""
加载预训练权重到编码器
"""
print("Loading pretrained encoders...")
self.model = ArrayForceInversionModel(
num_duns=self.num_duns,
encoder_dim=32,
dropout_rate=0.2
).to(device)
try:
# 加载预训练权重
pretrained_state = torch.load(self.pretrained_path, map_location=device)
print(f"✓ Loaded pretrained model from {self.pretrained_path}")
# 提取MLP权重到所有编码器
mlp_state_dict = {}
for k, v in pretrained_state.items():
if k.startswith('mlp.'):
new_k = k.replace('mlp.', '', 1)
mlp_state_dict[new_k] = v
# 加载到所有墩墩编码器
if mlp_state_dict:
for dun_encoder in self.model.dun_encoders:
encoder_dict = dun_encoder.state_dict()
# 匹配并加载权重
for key in encoder_dict.keys():
if key in mlp_state_dict:
try:
encoder_dict[key].copy_(mlp_state_dict[key])
except Exception as e:
print(f" ✗ Failed to load {key}: {e}")
print("✓ Pretrained weights loaded to all dun encoders")
# 设置归一化参数
if 'x_mean' in pretrained_state:
B_mean = pretrained_state['x_mean'].reshape(-1)
B_std = pretrained_state['x_std'].reshape(-1)
F_mean = pretrained_state['y_mean'].reshape(-1)
F_std = pretrained_state['y_std'].reshape(-1)
self.model.set_normalization_params(B_mean, B_std, F_mean, F_std)
print("✓ Normalization parameters set")
except Exception as e:
print(f"⚠ Failed to load pretrained weights: {e}")
print(" Using random initialization")
# 统计参数
total_params = sum(p.numel() for p in self.model.parameters())
encoder_params = sum(p.numel() for p in self.model.dun_encoders.parameters())
interaction_params = sum(p.numel() for p in self.model.interaction_network.parameters())
predictor_params = sum(p.numel() for p in self.model.force_predictor.parameters())
print(f"\nModel Architecture:")
print(f" Total Parameters: {total_params:,}")
print(f" Encoder Parameters: {encoder_params:,} ({self.num_duns} × {encoder_params//self.num_duns:,})")
print(f" Interaction Parameters: {interaction_params:,}")
print(f" Predictor Parameters: {predictor_params:,}")
return self.model
def progressive_fine_tune(self, B_array, F_total, B_source=None, F_source=None,
val_split=0.15, batch_size=64):
"""
两阶段渐进式微调
"""
# 数据分割
n_samples = B_array.size(0)
indices = torch.randperm(n_samples)
val_size = int(val_split * n_samples)
train_idx, val_idx = indices[val_size:], indices[:val_size]
B_train = B_array[train_idx]
F_train = F_total[train_idx]
B_val = B_array[val_idx]
F_val = F_total[val_idx]
print(f"Data split: Train={len(train_idx)}, Val={len(val_idx)}")
# Stage 1: 冻结编码器，训练交互网络和预测器
print("\n" + "="*70)
print("STAGE 1: Domain Adaptation (Frozen Encoders)")
print("="*70)
# 冻结编码器
for encoder in self.model.dun_encoders:
for param in encoder.parameters():
param.requires_grad = False
trainable_params = list(self.model.interaction_network.parameters()) + \
list(self.model.force_predictor.parameters()) + \
[self.model.position_encoding]
stage1_history = self._train_stage(
B_train, F_train, B_val, F_val,
B_source, F_source,
trainable_params=trainable_params,
epochs=100, lr=1e-4, batch_size=batch_size,
stage_name="Stage1"
)
self.history.append({'stage': 'Stage1', 'history': stage1_history})
# Stage 2: 整体微调
print("\n" + "="*70)
print("STAGE 2: End-to-End Fine-tuning")
print("="*70)
# 解冻所有参数
for param in self.model.parameters():
param.requires_grad = True
stage2_history = self._train_stage(
B_train, F_train, B_val, F_val,
None, None, # 第二阶段不使用源域数据
trainable_params=self.model.parameters(),
epochs=150, lr=5e-5, batch_size=batch_size,
stage_name="Stage2"
)
self.history.append({'stage': 'Stage2', 'history': stage2_history})
# 保存最终模型
torch.save(self.model.state_dict(), os.path.join(self.output_folder, 'model_final.pth'))
print("✓ Final model saved")
self._plot_training_history()
def _train_stage(self, B_train, F_train, B_val, F_val,
B_source, F_source, trainable_params,
epochs, lr, batch_size, stage_name):
"""
训练单个阶段
"""
optimizer = optim.Adam(trainable_params, lr=lr, weight_decay=1e-5)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
# 数据加载器
train_dataset = torch.utils.data.TensorDataset(B_train, F_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_dataset = torch.utils.data.TensorDataset(B_val, F_val)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
# 源域数据加载器（用于域适应）
source_loader = None
if B_source is not None and F_source is not None:
source_dataset = torch.utils.data.TensorDataset(B_source, F_source)
source_loader = torch.utils.data.DataLoader(source_dataset, batch_size=batch_size, shuffle=True)
history = {
'train_loss': [], 'train_task_loss': [], 'train_mmd_loss': [], 'train_physics_loss': [],
'val_loss': [], 'val_metrics': [], 'learning_rate': []
}
best_val_loss = float('inf')
patience, patience_counter = 20, 0
for epoch in range(epochs):
# 训练
self.model.train()
epoch_losses = {k: 0.0 for k in ['total', 'task', 'mmd', 'physics']}
for batch_idx, (B_batch, F_batch) in enumerate(train_loader):
optimizer.zero_grad()
# 域适应：从源域采样批次
if source_loader is not None:
try:
source_iter = source_loader.__iter__()
B_source_batch, F_source_batch = next(source_iter)
except:
source_iter = source_loader.__iter__()
B_source_batch, F_source_batch = next(source_iter)
# 对齐批次大小
min_size = min(B_batch.size(0), B_source_batch.size(0))
B_batch = B_batch[:min_size]
F_batch = F_batch[:min_size]
B_source_batch = B_source_batch[:min_size]
# 前向传播
if source_loader is not None:
# 域适应模式
F_pred, target_features, _ = self.model(B_batch, return_features=True)
_, source_features, _ = self.model(B_source_batch, return_features=True)
# 计算各种损失
task_loss = self.mse_loss(F_pred, F_batch)
# MMD损失
mmd_loss = 0.0
for i in range(self.num_duns):
mmd_loss += self.mmd_loss(source_features[i], target_features[i])
mmd_loss = mmd_loss / self.num_duns
# 物理一致性损失
physics_loss = self.physics_loss(target_features)
# 总损失
total_loss = task_loss + 0.3 * mmd_loss + 0.1 * physics_loss
else:
# 标准训练模式
F_pred = self.model(B_batch)
task_loss = self.mse_loss(F_pred, F_batch)
total_loss = task_loss
mmd_loss = torch.tensor(0.0)
physics_loss = torch.tensor(0.0)
total_loss.backward()
torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0)
optimizer.step()
# 记录损失
epoch_losses['total'] += total_loss.item()
epoch_losses['task'] += task_loss.item()
epoch_losses['mmd'] += mmd_loss.item() if isinstance(mmd_loss, torch.Tensor) else mmd_loss
epoch_losses['physics'] += physics_loss.item() if isinstance(physics_loss, torch.Tensor) else physics_loss
# 记录平均损失
num_batches = len(train_loader)
history['train_loss'].append(epoch_losses['total'] / num_batches)
history['train_task_loss'].append(epoch_losses['task'] / num_batches)
history['train_mmd_loss'].append(epoch_losses['mmd'] / num_batches)
history['train_physics_loss'].append(epoch_losses['physics'] / num_batches)
history['learning_rate'].append(optimizer.param_groups[0]['lr'])
# 验证
val_loss, val_metrics = self._validate(val_loader)
history['val_loss'].append(val_loss)
history['val_metrics'].append(val_metrics)
scheduler.step()
# 打印进度
if (epoch + 1) % 10 == 0 or epoch == 0:
print(f"{stage_name} Epoch {epoch+1}/{epochs}:")
print(f" Train Loss: {history['train_loss'][-1]:.6f}")
print(f" Val Loss: {val_loss:.6f}")
print(f" Force R²: {val_metrics['overall']['r2']:.4f}, MAE: {val_metrics['overall']['mae']:.4f}")
if source_loader is not None:
print(f" MMD Loss: {history['train_mmd_loss'][-1]:.6f}")
# 早停
if val_loss < best_val_loss:
best_val_loss = val_loss
patience_counter = 0
# 保存最佳模型
torch.save(self.model.state_dict(),
os.path.join(self.output_folder, f'model_{stage_name}_best.pth'))
else:
patience_counter += 1
if patience_counter >= patience:
print(f"Early stopping at epoch {epoch+1}")
# 加载最佳模型
self.model.load_state_dict(torch.load(
os.path.join(self.output_folder, f'model_{stage_name}_best.pth')))
break
return history
def _validate(self, val_loader):
"""验证过程"""
self.model.eval()
val_loss = 0.0
all_preds, all_targets = [], []
with torch.no_grad():
for B_batch, F_batch in val_loader:
F_pred = self.model(B_batch)
loss = self.mse_loss(F_pred, F_batch)
val_loss += loss.item()
all_preds.append(F_pred.cpu())
all_targets.append(F_batch.cpu())
# 计算指标
all_preds = torch.cat(all_preds, dim=0)
all_targets = torch.cat(all_targets, dim=0)
metrics = compute_metrics(all_preds, all_targets)
return val_loss / len(val_loader), metrics
def evaluate(self, B_test, F_test, plot_name="final_evaluation"):
"""最终评估"""
self.model.eval()
with torch.no_grad():
F_pred = self.model(B_test)
metrics = compute_metrics(F_pred, F_test)
print(f"\n{'='*70}")
print("FINAL EVALUATION RESULTS")
print(f"{'='*70}")
for comp in ['Fx', 'Fy', 'Fz']:
print(f"{comp}: R² = {metrics[comp]['r2']:.6f}, MAE = {metrics[comp]['mae']:.6f}")
print(f"Overall: R² = {metrics['overall']['r2']:.6f}, MAE = {metrics['overall']['mae']:.6f}")
# 绘制结果
self._plot_results(F_pred.cpu().numpy(), F_test.cpu().numpy(),
os.path.join(self.output_folder, f"{plot_name}.png"))
return metrics
def _plot_results(self, preds, targets, save_path):
"""绘制预测结果"""
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
components = ['Fx', 'Fy', 'Fz']
for i, (ax, comp) in enumerate(zip(axes, components)):
true_vals = targets[:, i]
pred_vals = preds[:, i]
ax.scatter(true_vals, pred_vals, alpha=0.6, s=20)
min_val, max_val = min(true_vals.min(), pred_vals.min()), max(true_vals.max(), pred_vals.max())
ax.plot([min_val, max_val], [min_val, max_val], 'r--', alpha=0.8)
r2 = r2_score(true_vals, pred_vals)
mae = mean_absolute_error(true_vals, pred_vals)
ax.set_xlabel('True Value')
ax.set_ylabel('Predicted Value')
ax.set_title(f'{comp}\nR² = {r2:.4f}, MAE = {mae:.4f}')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
print(f"✓ Evaluation plot saved: {save_path}")
def _plot_training_history(self):
"""绘制训练历史"""
fig, axes = plt.subplots(2, 2, figsize=(15, 10))
# 准备数据
all_train_loss, all_val_loss = [], []
all_val_r2 = []
stage_boundaries = [0]
for stage_info in self.history:
hist = stage_info['history']
all_train_loss.extend(hist['train_loss'])
all_val_loss.extend(hist['val_loss'])
all_val_r2.extend([m['overall']['r2'] for m in hist['val_metrics']])
stage_boundaries.append(len(all_train_loss))
epochs = range(len(all_train_loss))
# 损失曲线
axes[0,0].plot(epochs, all_train_loss, label='Train Loss', linewidth=2)
axes[0,0].plot(epochs, all_val_loss, label='Val Loss', linewidth=2)
for boundary in stage_boundaries[1:-1]:
axes[0,0].axvline(x=boundary, color='red', linestyle='--', alpha=0.7, label='Stage Change')
axes[0,0].set_ylabel('Loss')
axes[0,0].set_title('Training and Validation Loss')
axes[0,0].legend()
axes[0,0].grid(True, alpha=0.3)
axes[0,0].set_yscale('log')
# R²曲线
axes[0,1].plot(epochs, all_val_r2, label='Force R²', color='green', linewidth=2)
for boundary in stage_boundaries[1:-1]:
axes[0,1].axvline(x=boundary, color='red', linestyle='--', alpha=0.7)
axes[0,1].set_ylabel('R² Score')
axes[0,1].set_title('Validation R² Score')
axes[0,1].legend()
axes[0,1].grid(True, alpha=0.3)
axes[0,1].set_ylim(0, 1)
# 损失分量（第一阶段）
if len(self.history) > 0:
stage1_hist = self.history[0]['history']
stages1_epochs = range(len(stage1_hist['train_loss']))
axes[1,0].plot(stages1_epochs, stage1_hist['train_task_loss'], label='Task Loss', linewidth=2)
axes[1,0].plot(stages1_epochs, stage1_hist['train_mmd_loss'], label='MMD Loss', linewidth=2)
axes[1,0].plot(stages1_epochs, stage1_hist['train_physics_loss'], label='Physics Loss', linewidth=2)
axes[1,0].set_xlabel('Epochs')
axes[1,0].set_ylabel('Loss')
axes[1,0].set_title('Stage 1: Loss Components')
axes[1,0].legend()
axes[1,0].grid(True, alpha=0.3)
axes[1,0].set_yscale('log')
# 最终结果对比
final_r2 = []
stage_names = []
for stage_info in self.history:
if stage_info['history']['val_metrics']:
final_r2.append(stage_info['history']['val_metrics'][-1]['overall']['r2'])
stage_names.append(stage_info['stage'])
axes[1,1].bar(range(len(final_r2)), final_r2, color=['blue', 'orange'], alpha=0.7)
axes[1,1].set_xticks(range(len(final_r2)))
axes[1,1].set_xticklabels(stage_names)
axes[1,1].set_ylabel('Final R² Score')
axes[1,1].set_title('Final Performance by Stage')
axes[1,1].grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig(os.path.join(self.output_folder, 'training_history.png'),
dpi=150, bbox_inches='tight')
plt.close()
print("✓ Training history plot saved")
###主函数###
def main():
# ===== 配置参数 =====
file_path = 'D:/Gen3/5_biaoding/figerpad_exp/IP03A/mat/大压头_all.mat'
train_data_path = 'D:/Gen3/5_biaoding/Data_processed/Dun3_F2B_Ftotal_processed.csv'
pretrained_path = 'D:/Gen3/4_B2F/MPL_Dun3_Ftotal_GymCenter_processed_force_GELU_64_32_1126/b_to_f_model_last.onnx'
output_folder = 'D:/Gen3/5_biaoding/ArrayForceInversion_Transformer_1126'
os.makedirs(output_folder, exist_ok=True)
# 数据加载
print("Loading data...")
B_array, F_total, num_duns = load_multi_dun_exp_data(file_path, device, num_samples=100000)
B_pretrain, F_pretrain = load_train_data(train_data_path, device)
# 创建单个墩墩的源域数据（用于域适应）
B_source = B_pretrain.unsqueeze(1) # (N, 1, 6, 3)
B_source = B_source.repeat(1, num_duns, 1, 1) # (N, num_duns, 6, 3)
F_source = F_pretrain
print(f"Source data: {B_source.shape}, {F_source.shape}")
# 初始化微调器
print("\nInitializing fine-tuner...")
tuner = ArrayForceInversionFineTuner(pretrained_path, output_folder, num_duns=num_duns)
tuner.load_pretrained_encoders(B_pretrain, F_pretrain)
# 执行两阶段微调
print("\nStarting progressive fine-tuning...")
tuner.progressive_fine_tune(
B_array, F_total,
B_source=B_source, F_source=F_source,
val_split=0.15, batch_size=64
)
# 最终评估
print("\nPerforming final evaluation...")
# 使用30%数据作为测试集
n_samples = B_array.size(0)
test_size = int(0.3 * n_samples)
test_indices = torch.randperm(n_samples)[:test_size]
B_test = B_array[test_indices]
F_test = F_total[test_indices]
metrics = tuner.evaluate(B_test, F_test)
print(f"\n✓ All results saved to: {output_folder}")
return tuner
if __name__ == "__main__":
tuner = main()
最新发布
磁力链接分组计数
扫一扫
377
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
