MLXID/DeepLearning
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix the loss calculation in 3-ph training

Browse Source
This commit is contained in:
xie_x1
2026-06-12 11:09:37 +02:00
parent 70f6acc6f8
commit 774c039f8c
2 changed files with 124 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Configs/train_3photon.yaml
+8 -7
View File
@@ -7,9 +7,9 @@ data:
energy: 12 ### in keV
batch_size_train: 4096
batch_size_val: 1024
batch_size_test: 1024
num_workers: 32
batch_size_val: 8192
batch_size_test: 8192
num_workers: 16
train_file_range: [0, 12]
val_file_range: [13, 14]
test_file_range: [15, 15]
@@ -17,15 +17,16 @@ data:
n_size: 9 ### size of sub-images containing 3 photons
model:
version: "260529"
version: "260611" # 260529
training:
epochs: 1000
learning_rate: 1.0e-3
weight_decay: 1.0e-4
weight_decay: 0
scheduler_factor: 0.7
scheduler_patience: 5
checkpoint_epochs: [10, 30, 50, 100, 150, 300, 500, 1000]
checkpoint_epochs: [10, 20, 50, 100, 150, 300, 500, 1000]
loss:
type: "three_point_set_loss_smooth_l1"
type: "three_point_set_loss_smooth_l1"
smooth_l1_beta: 0.01
Train_3Photon.py
+116 -25
View File
@@ -29,38 +29,98 @@ def prepare_output_folder(conf):
# find the next index for experiment name
exp_index = 0
while True:
exp_name = f'{date}_3ph_{conf.data.energy}keV_v{conf.model.version}_{exp_index:02d}'
if not Path(f'Results/{exp_name}').exists():
exp_name = f'{date}_{conf.data.energy}keV_v{conf.model.version}_{exp_index:02d}'
if not Path(f'Results/3ph/{exp_name}').exists():
break
exp_index += 1
Path(f'Results/{exp_name}').mkdir(parents=True, exist_ok=True)
Path(f'Results/{exp_name}/Models').mkdir(parents=True, exist_ok=True)
Path(f'Results/{exp_name}/Plots').mkdir(parents=True, exist_ok=True)
OmegaConf.save(conf, f'Results/{exp_name}/config.yaml')
Path(f'Results/3ph/{exp_name}').mkdir(parents=True, exist_ok=True)
Path(f'Results/3ph/{exp_name}/Models').mkdir(parents=True, exist_ok=True)
Path(f'Results/3ph/{exp_name}/Plots').mkdir(parents=True, exist_ok=True)
OmegaConf.save(conf, f'Results/3ph/{exp_name}/config.yaml')
return exp_name
def get_loss_function(conf):
if conf.loss.type == "three_point_set_loss_smooth_l1":
def three_point_set_loss_smooth_l1(pred_xy, gt_xy):
loss_fn = torch.nn.SmoothL1Loss(reduction='none')
loss_fn = torch.nn.SmoothL1Loss(reduction='none', beta=conf.loss.smooth_l1_beta)
p1, p2, p3 = pred_xy[:,0], pred_xy[:,1], pred_xy[:,2]
g1, g2, g3 = gt_xy[:,0], gt_xy[:,1], gt_xy[:,2]
c_a = loss_fn(p1, g1).sum(dim=-1) + loss_fn(p2, g2).sum(dim=-1) + loss_fn(p3, g3).sum(dim=-1)
c_b = loss_fn(p1, g2).sum(dim=-1) + loss_fn(p2, g1).sum(dim=-1) + loss_fn(p3, g3).sum(dim=-1)
c_c = loss_fn(p1, g3).sum(dim=-1) + loss_fn(p2, g2).sum(dim=-1) + loss_fn(p3, g1).sum(dim=-1)
c_d = loss_fn(p1, g1).sum(dim=-1) + loss_fn(p2, g3).sum(dim=-1) + loss_fn(p3, g2).sum(dim=-1)
c_e = loss_fn(p1, g2).sum(dim=-1) + loss_fn(p2, g3).sum(dim=-1) + loss_fn(p3, g1).sum(dim=-1)
c_f = loss_fn(p1, g3).sum(dim=-1) + loss_fn(p2, g1).sum(dim=-1) + loss_fn(p3, g2).sum(dim=-1)
return torch.minimum(torch.minimum(torch.minimum(torch.minimum(torch.minimum(c_a, c_b), c_c), c_d), c_e), c_f).mean()
# 1. 匹配阶段：计算全部 6 种排列的 L2 物理距离 (不参与梯度计算)
with torch.no_grad():
def calc_cost(ga, gb, gc):
return ((p1 - ga)**2).sum(dim=-1) + \
((p2 - gb)**2).sum(dim=-1) + \
((p3 - gc)**2).sum(dim=-1)
# 枚举 3个光子的 6 种全排列 (Permutations)
c0 = calc_cost(g1, g2, g3)
c1 = calc_cost(g1, g3, g2)
c2 = calc_cost(g2, g1, g3)
c3 = calc_cost(g2, g3, g1)
c4 = calc_cost(g3, g1, g2)
c5 = calc_cost(g3, g2, g1)
# 将 6 种 cost 堆叠起来，形状变为 [B, 6]
costs = torch.stack([c0, c1, c2, c3, c4, c5], dim=1)
# 找到每个 batch 中 cost 最小的那个排列的索引 (0 到 5)
min_indices = torch.argmin(costs, dim=1) # [B]
# 预定义好 6 种排列对应的物理索引 (0代表g1, 1代表g2, 2代表g3)
perms = torch.tensor([
[0, 1, 2], # c0
[0, 2, 1], # c1
[1, 0, 2], # c2
[1, 2, 0], # c3
[2, 0, 1], # c4
[2, 1, 0] # c5
], device=pred_xy.device)
# 获取当前 batch 最优的排列索引组合，形状为 [B, 3]
best_perms = perms[min_indices]
# 将 Ground Truth 堆叠为 [B, 3, 2] 的统一张量
g_all = torch.stack([g1, g2, g3], dim=1)
# 利用高级索引一次性完成重组
batch_idx = torch.arange(pred_xy.size(0), device=pred_xy.device).unsqueeze(-1)
g_matched = g_all[batch_idx, best_perms] # [B, 3, 2]
# 2. 惩罚阶段：对重组后完美对齐的坐标计算 L2 Loss
loss_p1 = loss_fn(p1, g_matched[:, 0]).sum(dim=-1)
loss_p2 = loss_fn(p2, g_matched[:, 1]).sum(dim=-1)
loss_p3 = loss_fn(p3, g_matched[:, 2]).sum(dim=-1)
loss = (loss_p1 + loss_p2 + loss_p3).mean() / 3.0
return loss
return three_point_set_loss_smooth_l1
def get_mse(pred_xy, gt_xy):
p1, p2, p3 = pred_xy[:,0], pred_xy[:,1], pred_xy[:,2]
g1, g2, g3 = gt_xy[:,0], gt_xy[:,1], gt_xy[:,2]
with torch.no_grad():
def calc_cost(ga, gb, gc):
return ((p1 - ga)**2).sum(dim=-1) + \
((p2 - gb)**2).sum(dim=-1) + \
((p3 - gc)**2).sum(dim=-1)
c0 = calc_cost(g1, g2, g3)
c1 = calc_cost(g1, g3, g2)
c2 = calc_cost(g2, g1, g3)
c3 = calc_cost(g2, g3, g1)
c4 = calc_cost(g3, g1, g2)
c5 = calc_cost(g3, g2, g1)
costs = torch.stack([c0, c1, c2, c3, c4, c5], dim=1)
min_costs, _ = torch.min(costs, dim=1) # [B]
mse_1d = min_costs.mean() / 6.0 # num_photons=3, num_axes=2
return mse_1d
def train(model, trainLoader, optimizer, loss_fn):
model.train()
batchLoss = 0
sum_squared_error_1d = 0; event_count = 0
for batch_idx, (sample, label) in enumerate(trainLoader):
sample, label = sample.cuda(), label.cuda()
x1, y1, z1, e1 = label[:, 0, 0], label[:, 0, 1], label[:, 0, 2], label[:, 0, 3]
@@ -74,13 +134,18 @@ def train(model, trainLoader, optimizer, loss_fn):
loss.backward()
optimizer.step()
batchLoss += loss.item() * sample.shape[0]
mse_1d = get_mse(pred_xy, gt_xy)
sum_squared_error_1d += mse_1d.item() * sample.shape[0]
event_count += sample.shape[0]
avgLoss = batchLoss / len(trainLoader.dataset) / 6 ### divide by 6 to get the average loss per photon per axis
print(f"[Train]\t Average Loss: {avgLoss:.6f} (RMS = {np.sqrt(avgLoss):.6f})")
return avgLoss
rms_1d = np.sqrt(sum_squared_error_1d / event_count)
print(f"[Train]\t Average Loss: {avgLoss:.6f} (RMS = {rms_1d:.6f})")
return avgLoss, rms_1d
def evaluate(model, valLoader, loss_fn):
model.eval()
batchLoss = 0
sum_squared_error_1d = 0; event_count = 0
with torch.no_grad():
for batch_idx, (sample, label) in enumerate(valLoader):
sample, label = sample.cuda(), label.cuda()
@@ -91,10 +156,14 @@ def evaluate(model, valLoader, loss_fn):
output = model(sample)
pred_xy = torch.stack((output[:,0:2], output[:,2:4], output[:,4:6]), axis=1)
loss = loss_fn(pred_xy, gt_xy)
mse_1d = get_mse(pred_xy, gt_xy)
sum_squared_error_1d += mse_1d.item() * sample.shape[0]
event_count += sample.shape[0]
batchLoss += loss.item() * sample.shape[0]
avgLoss = batchLoss / len(valLoader.dataset) / 6 ### divide by 6 to get the average loss per photon per axis
print(f"[Val]\t Average Loss: {avgLoss:.6f} (RMS = {np.sqrt(avgLoss):.6f})")
return avgLoss
rms_1d = np.sqrt(sum_squared_error_1d / event_count)
print(f"[Val]\t Average Loss: {avgLoss:.6f} (RMS = {rms_1d:.6f})")
return avgLoss, rms_1d
def get_dataloaders(conf):
"""construct all dataloaders"""
@@ -136,7 +205,23 @@ def plot_loss_curves(train_losses, val_losses, test_loss, exp_name, conf):
plt.legend()
plt.grid()
plotName = f'loss_curve_triplePhoton_{conf.model.version}.png'
plt.savefig(f'Results/{exp_name}/Plots/{plotName}')
plt.savefig(f'Results/3ph/{exp_name}/Plots/{plotName}')
plt.close()
def plot_rms_curves(train_rms, val_rms, test_rms, exp_name, conf):
import matplotlib.pyplot as plt
plt.figure(figsize=(8,6))
plt.plot(train_rms, label='Train RMS')
plt.plot(val_rms, label='Val RMS')
if test_rms > 0:
plt.axhline(y=test_rms, color='green', linestyle='--', label='Test RMS')
plt.xlabel('Epoch')
plt.ylabel('RMS Error [pixels]')
plt.yscale('log')
plt.legend()
plt.grid()
plotName = f'rms_curve_triplePhoton_{conf.model.version}.png'
plt.savefig(f'Results/3ph/{exp_name}/Plots/{plotName}')
plt.close()
def get_model_name(conf):
@@ -155,17 +240,23 @@ if __name__ == "__main__":
trainLoader, valLoader, testLoader = get_dataloaders(conf)
TrainLosses, ValLosses = [], []
train_rms, val_rms = [], []
for epoch in tqdm(range(1, conf.training.epochs + 1)):
train_loss = train(model, trainLoader, optimizer, loss_fn)
val_loss = evaluate(model, valLoader, loss_fn)
train_loss, train_rms_epoch = train(model, trainLoader, optimizer, loss_fn)
val_loss, val_rms_epoch = evaluate(model, valLoader, loss_fn)
TrainLosses.append(train_loss)
ValLosses.append(val_loss)
train_rms.append(train_rms_epoch)
val_rms.append(val_rms_epoch)
scheduler.step(val_loss)
print(f"Learning Rate: {optimizer.param_groups[0]['lr']:.2e}")
if epoch in conf.training.checkpoint_epochs or epoch == conf.training.epochs:
modelName = get_model_name(conf)
torch.save(model.state_dict(), f'Results/{exp_name}/Models/{modelName}_E{epoch}.pth')
torch.save(model.state_dict(), f'Results/3ph/{exp_name}/Models/{modelName}_E{epoch}.pth')
print(f"Saved model checkpoint: {modelName}_E{epoch}.pth")
plot_loss_curves(TrainLosses, ValLosses, test_loss=-1, exp_name=exp_name, conf=conf)
test_loss = evaluate(model, testLoader, loss_fn)
plot_loss_curves(TrainLosses, ValLosses, test_loss=test_loss, exp_name=exp_name, conf=conf)
plot_rms_curves(train_rms, val_rms, test_rms=-1, exp_name=exp_name, conf=conf)
test_loss, test_rms = evaluate(model, testLoader, loss_fn)
plot_loss_curves(TrainLosses, ValLosses, test_loss=test_loss, exp_name=exp_name, conf=conf)
plot_rms_curves(train_rms, val_rms, test_rms=test_rms, exp_name=exp_name, conf=conf)