167 lines
7.8 KiB
Python
167 lines
7.8 KiB
Python
import sys
|
|
sys.path.append('./src')
|
|
from omegaconf import OmegaConf ### for yaml config parsing
|
|
import torch
|
|
import numpy as np
|
|
import torch.optim as optim
|
|
from tqdm import tqdm
|
|
from torchinfo import summary
|
|
from pathlib import Path
|
|
|
|
from models import get_triple_photon_model_class
|
|
from datasets import triplePhotonDataset
|
|
|
|
### random seed for reproducibility
|
|
torch.manual_seed(0)
|
|
torch.cuda.manual_seed(0)
|
|
np.random.seed(0)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
|
|
conf = OmegaConf.load("Configs/train_3photon.yaml")
|
|
|
|
def prepare_output_folder(conf):
|
|
from datetime import datetime
|
|
date = datetime.now().strftime("%y%m%d") ## YYMMDD format
|
|
# 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():
|
|
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')
|
|
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')
|
|
|
|
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()
|
|
|
|
return three_point_set_loss_smooth_l1
|
|
|
|
def train(model, trainLoader, optimizer, loss_fn):
|
|
model.train()
|
|
batchLoss = 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]
|
|
x2, y2, z2, e2 = label[:, 1, 0], label[:, 1, 1], label[:, 1, 2], label[:, 1, 3]
|
|
x3, y3, z3, e3 = label[:, 2, 0], label[:, 2, 1], label[:, 2, 2], label[:, 2, 3]
|
|
gt_xy = torch.stack((torch.stack((x1, y1), axis=1), torch.stack((x2, y2), axis=1), torch.stack((x3, y3), axis=1)), axis=1)
|
|
optimizer.zero_grad()
|
|
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)
|
|
loss.backward()
|
|
optimizer.step()
|
|
batchLoss += loss.item() * 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
|
|
|
|
def evaluate(model, valLoader, loss_fn):
|
|
model.eval()
|
|
batchLoss = 0
|
|
with torch.no_grad():
|
|
for batch_idx, (sample, label) in enumerate(valLoader):
|
|
sample, label = sample.cuda(), label.cuda()
|
|
x1, y1, z1, e1 = label[:, 0, 0], label[:, 0, 1], label[:, 0, 2], label[:, 0, 3]
|
|
x2, y2, z2, e2 = label[:, 1, 0], label[:, 1, 1], label[:, 1, 2], label[:, 1, 3]
|
|
x3, y3, z3, e3 = label[:, 2, 0], label[:, 2, 1], label[:, 2, 2], label[:, 2, 3]
|
|
gt_xy = torch.stack((torch.stack((x1, y1), axis=1), torch.stack((x2, y2), axis=1), torch.stack((x3, y3), axis=1)), axis=1)
|
|
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)
|
|
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
|
|
|
|
def get_dataloaders(conf):
|
|
"""construct all dataloaders"""
|
|
datasets = {}
|
|
loaders = {}
|
|
splits = ['Train', 'Val', 'Test']
|
|
keys = ['train_files', 'val_files', 'test_files']
|
|
batch_keys = ['batch_size_train', 'batch_size_val', 'batch_size_test']
|
|
file_range_keys = ['train_file_range', 'val_file_range', 'test_file_range']
|
|
|
|
for split, key, batch_key, file_range_key in zip(splits, keys, batch_keys, file_range_keys):
|
|
files = [f"{conf.data.sample_folder}/pileupOf3phs_sample_{i}.npz" for i in range(conf.data[file_range_key][0], conf.data[file_range_key][1] + 1)]
|
|
|
|
datasets[split] = triplePhotonDataset(
|
|
files,
|
|
sampleRatio = conf.data.sample_ratio,
|
|
datasetName = split.capitalize(),
|
|
)
|
|
|
|
loaders[split] = torch.utils.data.DataLoader(
|
|
datasets[split],
|
|
batch_size=conf.data[batch_key],
|
|
shuffle=(split=='Train'),
|
|
num_workers=conf.data.num_workers,
|
|
pin_memory=True
|
|
)
|
|
return loaders['Train'], loaders['Val'], loaders['Test']
|
|
|
|
def plot_loss_curves(train_losses, val_losses, test_loss, exp_name, conf):
|
|
import matplotlib.pyplot as plt
|
|
plt.figure(figsize=(8,6))
|
|
plt.plot(train_losses, label='Train Loss')
|
|
plt.plot(val_losses, label='Val Loss')
|
|
if test_loss > 0:
|
|
plt.axhline(y=test_loss, color='green', linestyle='--', label='Test Loss')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('MSE Loss')
|
|
plt.yscale('log')
|
|
plt.legend()
|
|
plt.grid()
|
|
plotName = f'loss_curve_triplePhoton_{conf.model.version}.png'
|
|
plt.savefig(f'Results/{exp_name}/Plots/{plotName}')
|
|
|
|
def get_model_name(conf):
|
|
modelName = f'triplePhoton{conf.model.version}_{conf.data.energy}keV'
|
|
return modelName
|
|
|
|
if __name__ == "__main__":
|
|
exp_name = prepare_output_folder(conf)
|
|
model = get_triple_photon_model_class(conf.model.version)().cuda()
|
|
# summary(model, input_size=(128, 3, conf.data.n_size, conf.data.n_size))
|
|
|
|
loss_fn = get_loss_function(conf)
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=conf.training.learning_rate, weight_decay=conf.training.weight_decay)
|
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=conf.training.scheduler_factor, patience=conf.training.scheduler_patience)
|
|
|
|
trainLoader, valLoader, testLoader = get_dataloaders(conf)
|
|
|
|
TrainLosses, ValLosses = [], []
|
|
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)
|
|
TrainLosses.append(train_loss)
|
|
ValLosses.append(val_loss)
|
|
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')
|
|
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) |