Select Git revision
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
train_MADE.py 4.01 KiB
import torch
from torch import tensor as T
import numpy as np
from sklearn.model_selection import train_test_split
import h5py
import argparse
from distribution_estimation import models
from distribution_estimation import training
'''
Authors: Nikolas Frediani (Nikolas.Frediani@physik.lmu.de), Michael Krämer, Philipp Mertsch, Kathrin Nippel
Related publication: "SECRET: Stochasticity Emulator for Cosmic Ray Electrons", arXiv:2501.06011
'''
if __name__ == '__main__':
'''
This script trains MADE.
It loads and prreprocesses the dataset, initializes the model and trainer,
and trains the model for the specified number of epochs.
Usage: python train_MADE.py -e <number_of_epochs> -i <input_dimension> -l <hidden_layers> -o <number_of_mixture_components>
Arguments:
-f, --filename: The filename of the training dataset (defaults to the provided file)
-e, --n_epochs: Number of epochs to train for (required)
-s, --start_epoch: Epoch to start at. If 0, a new model is initialized. If >0, try to load an existing model and continue training from that epoch (default: 0)
-i, --input_dim: Input dimension (required)
-l, --layer: Adds a hidden layer with the specified number of nodes (required, can be repeated for multiple layers)
-o, --n_mixture: Number of mixture components in the model output. Used to calculate the output dimension (required)
-m, --masks: Number of masks used (default: 1)
-p, --model_path: Path where to look for saved models (default: 'models/')
Example usage: python train_MADE.py -e 500 -i 41 -l 100 -l 100 -o 10 -m 1
'''
parser = argparse.ArgumentParser() # example default run parameters: ...python ...train_made.py -e 500 -i 41 -l 100 -o 10 -m 1
parser.add_argument('-f','--filename',type=str,default='batches_regime_Klein-Nishina_SNrate_2.0e+04_dist_spiral_mode_full_causality_True_Ecut_1.0e+04_seeds_0-9999.hdf5',help='name of the file where data is stored')
parser.add_argument('-e','--n_epochs',type=int,required=True,help='number of epochs to train for')
parser.add_argument('-s','--start_epoch',type=int,default=0,help='epoch to start at. If 0, new model is initialised. If >0, try to load existing model and continue training from that epoch')
parser.add_argument('-i','--input_dim',type=int,required=True,help='input dimension')
parser.add_argument('-l','--layer',action='append',type=int,required=True,help='adds a hidden layers with _ number of nodes')
parser.add_argument('-o','--n_mixture',type=int,required=True,help='number of mixture components in the model output. Used to calculate the output dimension')
parser.add_argument('-m','--masks',type=int,default=1,help='number of masks used')
parser.add_argument('-p','--model_path',type=str,default='runs/models/',choices=['runs/models/','runs/temp/'],help='path where to look for saved models') #choices restriction may be removed if necessary
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# define energy bins
Ens = torch.logspace(1., 5., 41)[14:33]
# load data and standardise
f = h5py.File('datasets/'+args.filename, 'r')
print('dataset:',f['fluxes'])
fluxes = (Ens**3 * T(np.array(f['fluxes']),dtype=torch.float)[14:33].t() * 0.06).log10()
f.close()
norm_fluxes = (fluxes-fluxes.mean(dim=0))/fluxes.std(dim=0).to(device)
# prepare dataloaders
train_set,test_set = train_test_split(norm_fluxes,train_size=0.9,random_state=151)
train_data, val_data = torch.utils.data.DataLoader(train_set,50,shuffle=True), torch.utils.data.DataLoader(test_set,50,shuffle=True)
# initialise model and trainer
made = models.MADE(args.input_dim,args.layer,args.n_mixture,num_masks=args.masks).to(device)
my_trainer = training.Trainer(made, training.logL_loss, torch.optim.Adam(made.parameters()), train_data, val_data)
# load latest stage and train
if args.start_epoch > 0:
my_trainer.load_state(args.start_epoch,args.model_path)
my_trainer.train(args.n_epochs)